360 lines
11 KiB
Plaintext
360 lines
11 KiB
Plaintext
package service
|
|
|
|
import (
|
|
"context"
|
|
"io"
|
|
"mime/multipart"
|
|
"net/http"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"gorm.io/driver/sqlite"
|
|
"gorm.io/gorm"
|
|
"gorm.io/gorm/logger"
|
|
|
|
"github.com/gochat/gochat/internal/config"
|
|
"github.com/gochat/gochat/internal/model"
|
|
"github.com/gochat/gochat/internal/repository"
|
|
)
|
|
|
|
// --- MIME detection helper tests ---
|
|
|
|
func TestDetectUploadMIMEFromFilename(t *testing.T) {
|
|
tests := []struct {
|
|
filename string
|
|
want string
|
|
}{
|
|
{"photo.png", "image/png"},
|
|
{"photo.jpg", "image/jpeg"},
|
|
{"photo.jpeg", "image/jpeg"},
|
|
{"anim.gif", "image/gif"},
|
|
{"modern.webp", "image/webp"},
|
|
{"icon.svg", "image/svg+xml"},
|
|
{"song.mp3", "audio/mpeg"},
|
|
{"audio.ogg", "audio/ogg"},
|
|
{"sound.wav", "audio/wav"},
|
|
{"clip.webm", "audio/webm"},
|
|
{"video.mp4", "video/mp4"},
|
|
{"doc.pdf", "application/pdf"},
|
|
{"data.csv", "text/csv"},
|
|
{"readme.txt", "text/plain"},
|
|
{"sheet.xls", "application/vnd.ms-excel"},
|
|
{"sheet.xlsx", "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"},
|
|
{"letter.doc", "application/msword"},
|
|
{"letter.docx", "application/vnd.openxmlformats-officedocument.wordprocessingml.document"},
|
|
{"unknown.xyz", "application/octet-stream"},
|
|
{"noext", "application/octet-stream"},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.filename, func(t *testing.T) {
|
|
got := detectUploadMIMEFromFilename(tt.filename)
|
|
assert.Equal(t, tt.want, got)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestCategorizeUploadMIME(t *testing.T) {
|
|
tests := []struct {
|
|
mimeType string
|
|
want string
|
|
}{
|
|
{"image/png", "image"},
|
|
{"image/jpeg", "image"},
|
|
{"image/gif", "image"},
|
|
{"audio/mpeg", "audio"},
|
|
{"audio/ogg", "audio"},
|
|
{"video/mp4", "video"},
|
|
{"application/pdf", "file"},
|
|
{"text/csv", "file"},
|
|
{"text/plain", "file"},
|
|
{"application/octet-stream", ""}, // unsupported category
|
|
{"application/zip", ""}, // unsupported
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.mimeType, func(t *testing.T) {
|
|
got := categorizeUploadMIME(tt.mimeType)
|
|
assert.Equal(t, tt.want, got)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestIsUploadMIMEAllowed(t *testing.T) {
|
|
// Allowed MIME types per category
|
|
assert.True(t, isUploadMIMEAllowed("image", "image/png"))
|
|
assert.True(t, isUploadMIMEAllowed("image", "image/jpeg"))
|
|
assert.True(t, isUploadMIMEAllowed("image", "image/gif"))
|
|
assert.True(t, isUploadMIMEAllowed("audio", "audio/mpeg"))
|
|
assert.True(t, isUploadMIMEAllowed("audio", "audio/ogg"))
|
|
assert.True(t, isUploadMIMEAllowed("video", "video/mp4"))
|
|
assert.True(t, isUploadMIMEAllowed("file", "application/pdf"))
|
|
assert.True(t, isUploadMIMEAllowed("file", "text/csv"))
|
|
assert.True(t, isUploadMIMEAllowed("file", "application/octet-stream"))
|
|
|
|
// Disallowed MIME types
|
|
assert.False(t, isUploadMIMEAllowed("image", "audio/mpeg")) // wrong category
|
|
assert.False(t, isUploadMIMEAllowed("video", "image/png")) // wrong category
|
|
assert.False(t, isUploadMIMEAllowed("unknown", "image/png")) // unknown category
|
|
}
|
|
|
|
// --- Service integration tests using SQLite ---
|
|
|
|
func setupUploadServiceTest(t *testing.T) (*gorm.DB, *UploadService, string) {
|
|
t.Helper()
|
|
|
|
db, err := gorm.Open(sqlite.Open("file::memory:"), &gorm.Config{
|
|
Logger: logger.Default.LogMode(logger.Silent),
|
|
})
|
|
require.NoError(t, err, "failed to open test DB")
|
|
|
|
require.NoError(t, db.AutoMigrate(&model.DirectUpload{}), "failed to auto-migrate DirectUpload")
|
|
|
|
t.Cleanup(func() {
|
|
sqlDB, _ := db.DB()
|
|
sqlDB.Close()
|
|
})
|
|
|
|
// Create temp directory for file storage
|
|
tmpDir := t.TempDir()
|
|
|
|
cfg := &config.Config{
|
|
Storage: config.StorageConfig{
|
|
Provider: "local",
|
|
LocalPath: tmpDir,
|
|
MaxFileSize: 20 * 1024 * 1024, // 20MB
|
|
},
|
|
}
|
|
|
|
directUploadRepo := repository.NewDirectUploadRepo(db)
|
|
uploadService := NewUploadService(directUploadRepo, cfg)
|
|
|
|
return db, uploadService, tmpDir
|
|
}
|
|
|
|
// makeMultipartFileHeader creates a multipart.FileHeader for testing.
|
|
func makeMultipartFileHeader(filename string, contentType string, content string) *multipart.FileHeader {
|
|
// Create a temp file with the content
|
|
tmpFile, err := os.CreateTemp("", filename)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
tmpFile.WriteString(content)
|
|
tmpFile.Close()
|
|
|
|
// Build a multipart request to extract a FileHeader
|
|
body := strings.NewReader("")
|
|
req, _ := http.NewRequest("POST", "/", body)
|
|
|
|
// We can't easily create a multipart.FileHeader programmatically,
|
|
// so we use a real multipart request to create one.
|
|
// Alternative: use a simpler approach with the actual multipart writer.
|
|
|
|
// Since multipart.FileHeader is hard to mock, we create a real multipart form.
|
|
pr, pw := io.Pipe()
|
|
writer := multipart.NewWriter(pw)
|
|
|
|
go func() {
|
|
defer pw.Close()
|
|
part, err := writer.CreateFormFile("file", filename)
|
|
if err != nil {
|
|
return
|
|
}
|
|
part.Write([]byte(content))
|
|
writer.Close()
|
|
}()
|
|
|
|
req, err := http.NewRequest("POST", "/", pr)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
req.Header.Set("Content-Type", writer.FormDataContentType())
|
|
|
|
fileHeader, _, err := req.FormFile("file")
|
|
if err != nil {
|
|
// Fallback: construct a minimal FileHeader manually
|
|
// This is tricky, so let's use a different approach
|
|
panic(err)
|
|
}
|
|
|
|
return fileHeader
|
|
}
|
|
|
|
func TestUploadService_AccountUpload_Success(t *testing.T) {
|
|
db, svc, tmpDir := setupUploadServiceTest(t)
|
|
_ = db
|
|
_ = tmpDir
|
|
|
|
// Create a real multipart form to get a FileHeader
|
|
body := &strings.Builder{}
|
|
writer := multipart.NewWriter(body)
|
|
part, err := writer.CreateFormFile("file", "test_image.png")
|
|
require.NoError(t, err)
|
|
part.Write([]byte("fake png content"))
|
|
writer.Close()
|
|
|
|
req, err := http.NewRequest("POST", "/", strings.NewReader(body.String()))
|
|
require.NoError(t, err)
|
|
req.Header.Set("Content-Type", writer.FormDataContentType())
|
|
|
|
fileHeader, err := req.FormFile("file")
|
|
require.NoError(t, err)
|
|
|
|
result, svcErr := svc.AccountUpload(context.Background(), 1, AccountUploadRequest{
|
|
FileHeader: fileHeader,
|
|
})
|
|
require.NoError(t, svcErr, "AccountUpload should succeed")
|
|
|
|
assert.Equal(t, uint(1), result.UploadID)
|
|
assert.NotEmpty(t, result.UploadUUID)
|
|
assert.Equal(t, "test_image.png", result.OriginalName)
|
|
assert.Equal(t, "image", result.FileType)
|
|
assert.Equal(t, "image/png", result.MimeType)
|
|
assert.Equal(t, "pending", result.Status)
|
|
assert.NotZero(t, result.ExpiresAt)
|
|
|
|
// Verify file was saved on disk
|
|
files, _ := filepath.Glob(filepath.Join(tmpDir, "account", "1", "*.png"))
|
|
assert.GreaterOrEqual(t, len(files), 1, "file should be saved on disk")
|
|
}
|
|
|
|
func TestUploadService_AccountUpload_NoAccountID(t *testing.T) {
|
|
_, svc, _ := setupUploadServiceTest(t)
|
|
|
|
_, svcErr := svc.AccountUpload(context.Background(), 0, AccountUploadRequest{
|
|
FileHeader: nil,
|
|
})
|
|
assert.Error(t, svcErr)
|
|
assert.Contains(t, svcErr.Error(), "account_id is required")
|
|
}
|
|
|
|
func TestUploadService_AccountUpload_NoFile(t *testing.T) {
|
|
_, svc, _ := setupUploadServiceTest(t)
|
|
|
|
_, svcErr := svc.AccountUpload(context.Background(), 1, AccountUploadRequest{
|
|
FileHeader: nil,
|
|
})
|
|
assert.Error(t, svcErr)
|
|
assert.Contains(t, svcErr.Error(), "file is required")
|
|
}
|
|
|
|
func TestUploadService_WidgetDirectUpload_Success(t *testing.T) {
|
|
db, svc, tmpDir := setupUploadServiceTest(t)
|
|
_ = db
|
|
_ = tmpDir
|
|
|
|
// Create a real multipart form for widget upload
|
|
body := &strings.Builder{}
|
|
writer := multipart.NewWriter(body)
|
|
part, err := writer.CreateFormFile("file", "visitor_audio.mp3")
|
|
require.NoError(t, err)
|
|
part.Write([]byte("fake mp3 content"))
|
|
writer.Close()
|
|
|
|
req, err := http.NewRequest("POST", "/", strings.NewReader(body.String()))
|
|
require.NoError(t, err)
|
|
req.Header.Set("Content-Type", writer.FormDataContentType())
|
|
|
|
fileHeader, err := req.FormFile("file")
|
|
require.NoError(t, err)
|
|
|
|
result, svcErr := svc.WidgetDirectUpload(context.Background(), WidgetDirectUploadRequest{
|
|
FileHeader: fileHeader,
|
|
})
|
|
require.NoError(t, svcErr, "WidgetDirectUpload should succeed")
|
|
|
|
assert.NotEmpty(t, result.UploadUUID)
|
|
assert.Equal(t, "visitor_audio.mp3", result.OriginalName)
|
|
assert.Equal(t, "audio", result.FileType)
|
|
assert.Equal(t, "audio/mpeg", result.MimeType)
|
|
assert.Equal(t, "pending", result.Status)
|
|
|
|
// Widget uploads should go to "widget_direct" subdirectory
|
|
files, _ := filepath.Glob(filepath.Join(tmpDir, "widget_direct", "*.mp3"))
|
|
assert.GreaterOrEqual(t, len(files), 1, "widget file should be saved in widget_direct dir")
|
|
}
|
|
|
|
func TestUploadService_WidgetDirectUpload_NoFile(t *testing.T) {
|
|
_, svc, _ := setupUploadServiceTest(t)
|
|
|
|
_, svcErr := svc.WidgetDirectUpload(context.Background(), WidgetDirectUploadRequest{
|
|
FileHeader: nil,
|
|
})
|
|
assert.Error(t, svcErr)
|
|
assert.Contains(t, svcErr.Error(), "file is required")
|
|
}
|
|
|
|
func TestUploadService_AccountUpload_UnsupportedFileType(t *testing.T) {
|
|
_, svc, _ := setupUploadServiceTest(t)
|
|
|
|
// Create multipart form with unsupported file type
|
|
body := &strings.Builder{}
|
|
writer := multipart.NewWriter(body)
|
|
part, err := writer.CreateFormFile("file", "malware.exe")
|
|
require.NoError(t, err)
|
|
part.Write([]byte("fake exe content"))
|
|
writer.Close()
|
|
|
|
req, err := http.NewRequest("POST", "/", strings.NewReader(body.String()))
|
|
require.NoError(t, err)
|
|
req.Header.Set("Content-Type", writer.FormDataContentType())
|
|
|
|
fileHeader, err := req.FormFile("file")
|
|
require.NoError(t, err)
|
|
|
|
_, svcErr := svc.AccountUpload(context.Background(), 1, AccountUploadRequest{
|
|
FileHeader: fileHeader,
|
|
})
|
|
assert.Error(t, svcErr)
|
|
// exe files get "application/octet-stream" MIME which falls into "" (unsupported) category
|
|
assert.Contains(t, svcErr.Error(), "unsupported file type")
|
|
}
|
|
|
|
func TestUploadService_CleanupExpiredUploads(t *testing.T) {
|
|
db, svc, _ := setupUploadServiceTest(t)
|
|
|
|
// Insert an expired upload record directly
|
|
expiredUpload := &model.DirectUpload{
|
|
UploadUUID: "expired-uuid-123",
|
|
AccountID: 1,
|
|
Status: model.DirectUploadStatusExpired,
|
|
Source: model.DirectUploadSourceAccount,
|
|
OriginalName: "old_file.png",
|
|
FileType: "image",
|
|
MimeType: "image/png",
|
|
FileSize: 100,
|
|
FileURL: "/uploads/account/1/old_file.png",
|
|
ExpiresAt: time.Now().Add(-1 * time.Hour), // already expired
|
|
}
|
|
require.NoError(t, db.Create(expiredUpload).Error)
|
|
|
|
// Insert a non-expired upload
|
|
activeUpload := &model.DirectUpload{
|
|
UploadUUID: "active-uuid-456",
|
|
AccountID: 1,
|
|
Status: model.DirectUploadStatusPending,
|
|
Source: model.DirectUploadSourceAccount,
|
|
OriginalName: "new_file.png",
|
|
FileType: "image",
|
|
MimeType: "image/png",
|
|
FileSize: 200,
|
|
FileURL: "/uploads/account/1/new_file.png",
|
|
ExpiresAt: time.Now().Add(24 * time.Hour), // not expired yet
|
|
}
|
|
require.NoError(t, db.Create(activeUpload).Error)
|
|
|
|
count, err := svc.CleanupExpiredUploads(context.Background())
|
|
require.NoError(t, err)
|
|
assert.Equal(t, int64(1), count, "should clean up 1 expired upload")
|
|
|
|
// Verify active upload still exists
|
|
var remaining model.DirectUpload
|
|
require.NoError(t, db.Where("upload_uuid = ?", "active-uuid-456").First(&remaining).Error)
|
|
assert.Equal(t, model.DirectUploadStatusPending, remaining.Status)
|
|
}
|