Files
gochat/internal/service/upload_service_test.go_BAK
T
2026-06-04 15:44:48 +08:00

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)
}