Agent-System/internal/db/migrate_test.go

package db

import (
	"context"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/jackc/pgx/v5/pgxpool"
	"github.com/pressly/goose/v3"
)

func TestMigrationGate_ExpectedFilesMatchAppliedVersion(t *testing.T) {
	migDir := ResolveMigrationsDir()
	expected, err := ExpectedMigrationCount(migDir)
	if err != nil {
		t.Fatalf("count migrations: %v", err)
	}
	if expected < 1 {
		t.Fatal("expected at least one migration file")
	}

	dsn := os.Getenv("TEST_DATABASE_URL")
	if dsn == "" {
		t.Skip("TEST_DATABASE_URL not set; migration gate file count only")
	}

	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	pool, err := pgxpool.New(ctx, dsn)
	if err != nil {
		t.Skipf("postgres unavailable: %v", err)
	}
	defer pool.Close()
	if err := pool.Ping(ctx); err != nil {
		t.Skipf("postgres ping failed: %v", err)
	}
	if err := RunMigrations(ctx, pool, migDir); err != nil {
		t.Fatalf("RunMigrations: %v", err)
	}
	version, err := CurrentVersion(ctx, pool)
	if err != nil {
		t.Fatalf("CurrentVersion: %v", err)
	}
	if int(version) < expected {
		t.Fatalf("applied schema version %d behind migration files %d", version, expected)
	}
}

func TestMigration032_AllowsMigrationRepairKinds(t *testing.T) {
	files, err := ListMigrationFiles(ResolveMigrationsDir())
	if err != nil {
		t.Fatal(err)
	}
	var found string
	for _, name := range files {
		if strings.HasPrefix(name, "00032_") && strings.HasSuffix(name, ".sql") {
			found = name
			break
		}
	}
	if found == "" {
		t.Fatal("expected goose 00032 for migration_repair check constraints")
	}
	raw, err := os.ReadFile(filepath.Join(ResolveMigrationsDir(), found))
	if err != nil {
		t.Fatal(err)
	}
	text := string(raw)
	if !strings.Contains(text, "-- +goose Up") || !strings.Contains(text, "-- +goose Down") {
		t.Fatal("00032 must have Up and Down")
	}
	if !strings.Contains(text, "chk_snapshot_source_kind") || !strings.Contains(text, "chk_change_kind") {
		t.Fatal("00032 must replace source/change kind checks")
	}
	if strings.Count(text, "'migration_repair'") < 2 {
		t.Fatalf("00032 must allow migration_repair on both constraints, got %s", found)
	}
}

func TestGooseU9_00038_HasUpAndDown(t *testing.T) {
	files, err := ListMigrationFiles(ResolveMigrationsDir())
	if err != nil {
		t.Fatal(err)
	}
	var found string
	for _, name := range files {
		if strings.HasPrefix(name, "00038_") && strings.HasSuffix(name, ".sql") {
			found = name
			break
		}
	}
	if found == "" {
		t.Fatal("expected goose 00038 for dataset lifecycle")
	}
	raw, err := os.ReadFile(filepath.Join(ResolveMigrationsDir(), found))
	if err != nil {
		t.Fatal(err)
	}
	text := string(raw)
	if !strings.Contains(text, "-- +goose Up") || !strings.Contains(text, "-- +goose Down") {
		t.Fatal("00038 must have Up and Down")
	}
	for _, table := range []string{"dataset_registry", "dataset_backfill_jobs", "dataset_backfill_checkpoints", "dataset_backfill_queue", "dataset_destroy_requests"} {
		if !strings.Contains(text, table) {
			t.Fatalf("00038 must mention %s", table)
		}
	}
}

func TestGooseM2_00034to00036_DownThenUp(t *testing.T) {
	if testing.Short() {
		t.Skip("goose down is not part of -short")
	}
	dsn := SkipIfNoTestDatabaseURL(t)
	if strings.Contains(dsn, "5433") || strings.Contains(dsn, "agent_system_dev") ||
		(strings.Contains(dsn, "5434") && strings.Contains(dsn, "agent_system_test")) ||
		strings.Contains(dsn, "5435") {
		t.Fatal("TEST_DATABASE_URL points at a shared env port; refuse agent_system_dev/test/main")
	}
	migDir := ResolveMigrationsDir()
	files, err := ListMigrationFiles(migDir)
	if err != nil {
		t.Fatal(err)
	}
	want := []string{
		"00034_2026-09-14-workspace-lease-fencing.sql",
		"00035_2026-09-14-durable-event-delivery.sql",
		"00036_2026-09-14-final-artifact-vault.sql",
		"00037_2026-09-15-task-run-manifests.sql",
		"00038_2026-09-15-dataset-lifecycle.sql",
	}
	joined := strings.Join(files, "\n")
	for _, name := range want {
		if !strings.Contains(joined, name) {
			t.Fatalf("missing %s", name)
		}
	}

	ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
	defer cancel()
	admin, err := pgxpool.New(ctx, dsn)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(admin.Close)

	const isolated = "agent_rel_m2_goose"
	_, _ = admin.Exec(ctx, `DROP DATABASE IF EXISTS `+isolated)
	if _, err := admin.Exec(ctx, `CREATE DATABASE `+isolated); err != nil {
		t.Fatalf("create isolated db: %v", err)
	}
	t.Cleanup(func() {
		_, _ = admin.Exec(context.Background(), `DROP DATABASE IF EXISTS `+isolated)
	})

	m2dsn := rewriteDatabaseName(dsn, isolated)
	pool, err := pgxpool.New(ctx, m2dsn)
	if err != nil {
		t.Fatal(err)
	}
	defer pool.Close()
	if err := RunMigrations(ctx, pool, migDir); err != nil {
		t.Fatalf("up: %v", err)
	}
	ver, err := CurrentVersion(ctx, pool)
	if err != nil || ver < 38 {
		t.Fatalf("version after up: %d %v", ver, err)
	}

	sqlDB := OpenSQLDB(pool)
	defer sqlDB.Close()
	if err := goose.SetDialect("postgres"); err != nil {
		t.Fatal(err)
	}
	goose.SetTableName(schemaTable)
	if err := goose.DownToContext(ctx, sqlDB, migDir, 33); err != nil {
		t.Fatalf("down to 33: %v", err)
	}
	for _, table := range []string{"event_envelopes", "event_deliveries", "event_dead_letters", "final_artifacts", "task_run_manifests", "dataset_registry"} {
		var n int
		if err := pool.QueryRow(ctx, `SELECT COUNT(*) FROM information_schema.tables WHERE table_name=$1`, table).Scan(&n); err != nil {
			t.Fatal(err)
		}
		if n != 0 {
			t.Fatalf("%s still present after down", table)
		}
	}
	var col int
	if err := pool.QueryRow(ctx, `
		SELECT COUNT(*) FROM information_schema.columns
		WHERE table_name='document_records' AND column_name='lease_fencing_token'
	`).Scan(&col); err != nil {
		t.Fatal(err)
	}
	if col != 0 {
		t.Fatal("lease_fencing_token must drop with 00034 Down")
	}
	if err := RunMigrations(ctx, pool, migDir); err != nil {
		t.Fatalf("up after down: %v", err)
	}
}

func rewriteDatabaseName(dsn, name string) string {
	if u, err := url.Parse(dsn); err == nil && u.Scheme != "" {
		u.Path = "/" + name
		return u.String()
	}
	return dsn
}