gochat/backend/internal/database/upload_migration_test.go

package database

import (
	"bytes"
	"context"
	"fmt"
	"mime/multipart"
	"net/http/httptest"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gorm.io/driver/postgres"
	"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"
	"github.com/gochat/gochat/internal/service"
)

func openUploadMigrationPostgres(t *testing.T) (*gorm.DB, string) {
	t.Helper()
	if os.Getenv("GOCHAT_TEST_DB") == "sqlite" {
		t.Skip("requires PostgreSQL migration semantics")
	}
	dsn := os.Getenv("GOCHAT_TEST_DB_URL")
	if dsn == "" {
		dsn = "postgres://postgres:postgres@localhost:5432/gochat_test?sslmode=disable"
	}
	admin, err := gorm.Open(postgres.Open(dsn), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)})
	require.NoError(t, err)
	adminDB, err := admin.DB()
	require.NoError(t, err)
	t.Cleanup(func() { _ = adminDB.Close() })

	schema := fmt.Sprintf("upload_migration_%d", time.Now().UnixNano())
	require.NoError(t, admin.Exec("CREATE SCHEMA "+schema).Error)
	t.Cleanup(func() { _ = admin.Exec("DROP SCHEMA " + schema + " CASCADE").Error })

	migrationURL, err := url.Parse(dsn)
	require.NoError(t, err)
	require.NotEmpty(t, migrationURL.Scheme, "GOCHAT_TEST_DB_URL must be a PostgreSQL URL")
	query := migrationURL.Query()
	query.Set("search_path", schema+",public")
	migrationURL.RawQuery = query.Encode()

	db, err := gorm.Open(postgres.Open(migrationURL.String()), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)})
	require.NoError(t, err)
	sqlDB, err := db.DB()
	require.NoError(t, err)
	sqlDB.SetMaxOpenConns(1)
	t.Cleanup(func() { _ = sqlDB.Close() })
	return db, migrationURL.String()
}

func productionMigrationsPath(t *testing.T) string {
	t.Helper()
	path, err := filepath.Abs(filepath.Join("..", "..", "migrations"))
	require.NoError(t, err)
	return path
}

func TestUploadPostgresProductionMigrationsFromEmptySchema(t *testing.T) {
	db, dbURL := openUploadMigrationPostgres(t)
	require.NoError(t, RunMigrations(dbURL, productionMigrationsPath(t)))
	version, dirty, err := CurrentVersion(dbURL, productionMigrationsPath(t))
	require.NoError(t, err)
	assert.Equal(t, uint(86), version)
	assert.False(t, dirty)
	exerciseUploadProductionSchema(t, db)
}

func TestUploadPostgresProductionMigrationUpgradesLegacySchema(t *testing.T) {
	db, dbURL := openUploadMigrationPostgres(t)
	migrations := productionMigrationsPath(t)
	require.NoError(t, MigrateSteps(dbURL, migrations, 82))
	var legacyAccountID, legacyInboxID uint
	require.NoError(t, db.Raw("INSERT INTO accounts(name) VALUES ('legacy account') RETURNING id").Scan(&legacyAccountID).Error)
	require.NoError(t, db.Raw(`INSERT INTO inboxes(account_id, name, channel_type, channel_id)
		VALUES (?, 'legacy inbox', 'web_widget', 1) RETURNING id`, legacyAccountID).Scan(&legacyInboxID).Error)
	require.NoError(t, db.Exec(`INSERT INTO direct_uploads(id, account_id, file_name, file_url, file_size, content_type)
		VALUES (44, ?, 'legacy.png', '/uploads/account/1/legacy.png', 12, 'image/png')`, legacyAccountID).Error)
	require.NoError(t, db.Exec(`INSERT INTO widget_file_uploads(id, inbox_id, enabled, max_file_size, allowed_types)
		VALUES (7, ?, false, 4096, '["image/png"]')`, legacyInboxID).Error)

	require.NoError(t, MigrateSteps(dbURL, migrations, 1))
	assertLegacyUploadMapping(t, db)
	exerciseUploadProductionSchema(t, db)

	require.NoError(t, MigrateSteps(dbURL, migrations, -1))
	assert.True(t, db.Migrator().HasColumn("direct_uploads", "file_name"))
	assert.False(t, db.Migrator().HasColumn("direct_uploads", "upload_uuid"))
	var legacyName, legacyMIME string
	require.NoError(t, db.Raw("SELECT file_name, content_type FROM direct_uploads WHERE id = 44").Row().Scan(&legacyName, &legacyMIME))
	assert.Equal(t, "legacy.png", legacyName)
	assert.Equal(t, "image/png", legacyMIME)
	assertLegacyWidgetConfig(t, db, "widget_file_uploads")

	require.NoError(t, MigrateSteps(dbURL, migrations, 1))
	assertLegacyUploadMapping(t, db)
}

func assertLegacyUploadMapping(t *testing.T, db *gorm.DB) {
	t.Helper()
	var upload model.DirectUpload
	require.NoError(t, db.First(&upload, 44).Error)
	assert.NotEmpty(t, upload.UploadUUID)
	assert.Equal(t, model.DirectUploadStatusCompleted, upload.Status)
	assert.Equal(t, model.DirectUploadSourceAccount, upload.Source)
	assert.Equal(t, "legacy.png", upload.OriginalName)
	assert.Equal(t, "image", upload.FileType)
	assert.Equal(t, "image/png", upload.MimeType)
	assert.Equal(t, int64(12), upload.FileSize)
	assert.JSONEq(t, `{}`, string(upload.Metadata))
	assert.False(t, upload.ExpiresAt.IsZero())
	assertLegacyWidgetConfig(t, db, "widget_file_upload_configs")
}

func assertLegacyWidgetConfig(t *testing.T, db *gorm.DB, table string) {
	t.Helper()
	var enabled bool
	var maxFileSize int
	var allowedTypes string
	require.NoError(t, db.Raw("SELECT enabled, max_file_size, allowed_types::text FROM "+table+" WHERE id = 7").Row().Scan(&enabled, &maxFileSize, &allowedTypes))
	assert.False(t, enabled)
	assert.Equal(t, 4096, maxFileSize)
	assert.JSONEq(t, `["image/png"]`, allowedTypes)
}

func exerciseUploadProductionSchema(t *testing.T, db *gorm.DB) {
	t.Helper()
	ctx := context.Background()
	var accountID, inboxID, contactID, contactInboxID uint
	require.NoError(t, db.Raw("INSERT INTO accounts(name) VALUES (?) RETURNING id", "upload migration account").Scan(&accountID).Error)
	require.NoError(t, db.Raw(`INSERT INTO inboxes(account_id, name, channel_type, channel_id, enabled, channel_config)
		VALUES (?, ?, 'web_widget', 1, true, ?) RETURNING id`, accountID, "upload migration inbox", `{"website_token":"migration-website"}`).Scan(&inboxID).Error)
	require.NoError(t, db.Raw("INSERT INTO contacts(account_id, name) VALUES (?, ?) RETURNING id", accountID, "migration contact").Scan(&contactID).Error)
	require.NoError(t, db.Raw(`INSERT INTO contact_inboxes(contact_id, inbox_id, pubsub_token)
		VALUES (?, ?, ?) RETURNING id`, contactID, inboxID, "migration-widget").Scan(&contactInboxID).Error)
	require.NotZero(t, contactInboxID)

	tmpDir := t.TempDir()
	uploadService := service.NewUploadService(repository.NewDirectUploadRepo(db), &config.Config{
		Storage: config.StorageConfig{LocalPath: tmpDir, MaxFileSize: 20 << 20},
	}).WithWidgetAuth(repository.NewInboxRepo(db), repository.NewContactInboxRepo(db)).WithAccessDB(db)
	direct, err := uploadService.AccountDirectUpload(ctx, accountID, service.AccountDirectUploadRequest{
		FileHeader: uploadMigrationFileHeader(t, "direct.png"),
	})
	require.NoError(t, err)
	assert.Equal(t, string(model.DirectUploadStatusPending), direct.Status)
	_, ok := uploadService.ResolveAuthorizedUpload(ctx, direct.FileURL, accountID, "")
	assert.True(t, ok)
	_, ok = uploadService.ResolveAuthorizedUpload(ctx, direct.FileURL, accountID+1, "")
	assert.False(t, ok)

	widgetService := service.NewWidgetService(
		repository.NewInboxRepo(db),
		repository.NewContactRepo(db),
		repository.NewContactInboxRepo(db),
		repository.NewConversationRepo(db),
		repository.NewMessageRepo(db),
		nil, nil, nil,
		repository.NewWidgetFileUploadRepo(db),
		nil, nil, nil, nil,
	)
	widgetHeader := uploadMigrationFileHeader(t, "widget.png")
	widgetReader, err := widgetHeader.Open()
	require.NoError(t, err)
	t.Cleanup(func() { _ = widgetReader.Close() })
	widget, err := widgetService.StageFileUpload(ctx, service.WidgetUploadRequest{
		WebsiteToken: "migration-website",
		WidgetToken:  "migration-widget",
		FileHeader:   widgetHeader,
	}, widgetReader)
	require.NoError(t, err)
	status, err := widgetService.GetFileUploadStatus(ctx, "migration-website", widget.UploadUUID, "migration-widget")
	require.NoError(t, err)
	require.NotNil(t, status)
	assert.Equal(t, widget.UploadUUID, status.UploadUUID)

	widgetPath := filepath.Join(tmpDir, filepath.FromSlash(strings.TrimPrefix(widget.FileURL, "/uploads/")))
	require.NoError(t, os.MkdirAll(filepath.Dir(widgetPath), 0o755))
	require.NoError(t, os.WriteFile(widgetPath, uploadMigrationPNG, 0o600))
	_, ok = uploadService.ResolveAuthorizedUpload(ctx, widget.FileURL, 0, "migration-widget")
	assert.True(t, ok)
	_, ok = uploadService.ResolveAuthorizedUpload(ctx, widget.FileURL, 0, "wrong-widget")
	assert.False(t, ok)
}

var uploadMigrationPNG = []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n', 0, 0, 0, 0}

func uploadMigrationFileHeader(t *testing.T, name string) *multipart.FileHeader {
	t.Helper()
	var body bytes.Buffer
	writer := multipart.NewWriter(&body)
	part, err := writer.CreateFormFile("file", name)
	require.NoError(t, err)
	_, err = part.Write(uploadMigrationPNG)
	require.NoError(t, err)
	require.NoError(t, writer.Close())
	req := httptest.NewRequest("POST", "/", &body)
	req.Header.Set("Content-Type", writer.FormDataContentType())
	require.NoError(t, req.ParseMultipartForm(1<<20))
	t.Cleanup(func() { _ = req.MultipartForm.RemoveAll() })
	return req.MultipartForm.File["file"][0]
}