NianAIGC/backend/internal/migrations/migrations_integration_test.go

package migrations

import (
	"context"
	"os"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/jackc/pgx/v5"
)

// Requires an independently disposable database with no valuable contents.
// It never consults the application's DATABASE_URL or .env files.
func TestIntegrationRunOnDisposableDatabase(t *testing.T) {
	dsn := os.Getenv("NIAN_TEST_MIGRATION_DATABASE_URL")
	if dsn == "" {
		t.Skip("set NIAN_TEST_MIGRATION_DATABASE_URL for disposable PostgreSQL integration test")
	}
	ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
	defer cancel()
	conn, err := pgx.Connect(ctx, dsn)
	if err != nil {
		t.Fatal("connect disposable database:", err)
	}
	defer conn.Close(context.Background())
	if _, err := conn.Exec(ctx, `DO $$ BEGIN
		IF NOT EXISTS (SELECT 1 FROM pg_roles WHERE rolname = 'nian_migration_test') THEN
			CREATE ROLE nian_migration_test NOLOGIN;
		END IF;
	END $$`); err != nil {
		t.Fatal("create disposable test role:", err)
	}
	all, err := discover(sqlFiles)
	if err != nil {
		t.Fatal(err)
	}
	// Repair only markers this test might have left if an earlier assertion
	// stopped execution before normal cleanup.
	var trackingTable *string
	if err := conn.QueryRow(ctx, "SELECT to_regclass('public.schema_migrations')::text").Scan(&trackingTable); err != nil {
		t.Fatal(err)
	}
	if trackingTable != nil {
		if _, err := conn.Exec(ctx, "UPDATE schema_migrations SET checksum=$1 WHERE version=$2 AND checksum='test-mismatch'", all[0].checksum, all[0].version); err != nil {
			t.Fatal(err)
		}
		if _, err := conn.Exec(ctx, "DELETE FROM schema_migrations WHERE version='test_unknown'"); err != nil {
			t.Fatal(err)
		}
	}
	env := func(key string) string {
		switch key {
		case "DATABASE_AUTO_MIGRATE":
			return "true"
		case "ZHINIAN_DATA_BACKEND":
			return "postgres"
		case "DATABASE_MIGRATION_URL", "DATABASE_URL":
			return dsn
		case "DATABASE_APP_ROLE":
			return "nian_migration_test"
		case "NODE_ENV":
			return "production"
		default:
			return ""
		}
	}
	var wg sync.WaitGroup
	results := make(chan error, 2)
	for i := 0; i < 2; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			results <- Run(ctx, env)
		}()
	}
	wg.Wait()
	close(results)
	for err := range results {
		if err != nil {
			t.Fatal("concurrent migration:", err)
		}
	}
	if err := Run(ctx, env); err != nil {
		t.Fatal("repeat migration:", err)
	}
	var count int
	if err := conn.QueryRow(ctx, "SELECT count(*) FROM schema_migrations").Scan(&count); err != nil || count != len(all) {
		t.Fatalf("recorded migrations: count=%d, err=%v", count, err)
	}
	var privilege bool
	if err := conn.QueryRow(ctx, "SELECT has_table_privilege('nian_migration_test', 'public.payment_orders', 'SELECT,INSERT,UPDATE')").Scan(&privilege); err != nil || !privilege {
		t.Fatalf("payment order privileges: granted=%v, err=%v", privilege, err)
	}
	// A mismatch must prevent any later migration or startup.
	first := all[0]
	if _, err := conn.Exec(ctx, "UPDATE schema_migrations SET checksum='test-mismatch' WHERE version=$1", first.version); err != nil {
		t.Fatal(err)
	}
	defer func() {
		cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cleanupCancel()
		_, _ = conn.Exec(cleanupCtx, "UPDATE schema_migrations SET checksum=$1 WHERE version=$2", first.checksum, first.version)
		_, _ = conn.Exec(cleanupCtx, "DELETE FROM schema_migrations WHERE version='test_unknown'")
	}()
	if err := Run(ctx, env); err == nil || !strings.Contains(err.Error(), "checksum differs") {
		t.Fatalf("expected checksum rejection, got %v", err)
	}
	if _, err := conn.Exec(ctx, "UPDATE schema_migrations SET checksum=$1 WHERE version=$2", first.checksum, first.version); err != nil {
		t.Fatal(err)
	}
	if _, err := conn.Exec(ctx, "INSERT INTO schema_migrations(version, checksum) VALUES ('test_unknown', 'test')"); err != nil {
		t.Fatal(err)
	}
	if err := Run(ctx, env); err == nil || !strings.Contains(err.Error(), "absent from bundle") {
		t.Fatalf("expected unknown-version rejection, got %v", err)
	}
	if _, err := conn.Exec(ctx, "DELETE FROM schema_migrations WHERE version='test_unknown'"); err != nil {
		t.Fatal(err)
	}
	if _, err := conn.Exec(ctx, "SELECT pg_advisory_lock($1)", migrationLockID); err != nil {
		t.Fatal(err)
	}
	defer conn.Exec(context.Background(), "SELECT pg_advisory_unlock($1)", migrationLockID)
	shortCtx, shortCancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
	defer shortCancel()
	if err := Run(shortCtx, env); err == nil || !strings.Contains(err.Error(), "timed out") {
		t.Fatalf("expected advisory-lock timeout, got %v", err)
	}
	if _, err := conn.Exec(ctx, "SELECT pg_advisory_unlock($1)", migrationLockID); err != nil {
		t.Fatal(err)
	}
	// The outer BEGIN/COMMIT in a legacy migration must not commit its DDL
	// before the version/checksum row is inserted by apply.
	const syntheticVersion = "9999_atomicity_test"
	const syntheticTable = "public.nian_migration_atomicity_test"
	if _, err := conn.Exec(ctx, "DROP TABLE IF EXISTS "+syntheticTable); err != nil {
		t.Fatal(err)
	}
	if _, err := conn.Exec(ctx, "DELETE FROM schema_migrations WHERE version=$1", syntheticVersion); err != nil {
		t.Fatal(err)
	}
	synthetic := migration{
		version: syntheticVersion,
		sql:     "-- synthetic migration\nBEGIN;\nCREATE TABLE " + syntheticTable + " (id integer PRIMARY KEY);\nSELECT 1 / 0;\nCOMMIT;\n",
	}
	if err := apply(ctx, conn, append(all, synthetic), ""); err == nil || !strings.Contains(err.Error(), syntheticVersion) {
		t.Fatalf("expected synthetic migration failure with version, got %v", err)
	}
	var objectName *string
	if err := conn.QueryRow(ctx, "SELECT to_regclass($1)::text", syntheticTable).Scan(&objectName); err != nil {
		t.Fatal(err)
	}
	if objectName != nil {
		t.Fatalf("failed migration left DDL object %s", *objectName)
	}
	var recorded int
	if err := conn.QueryRow(ctx, "SELECT count(*) FROM schema_migrations WHERE version=$1", syntheticVersion).Scan(&recorded); err != nil || recorded != 0 {
		t.Fatalf("failed migration recorded version: count=%d, err=%v", recorded, err)
	}
}