easyai-ai-gateway/apps/api/cmd/migrate/main_test.go

package main

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

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgconn"
	"github.com/jackc/pgx/v5/pgxpool"
)

func TestSecurityEventSchemaMigrationsDefineCurrentLifecycle(t *testing.T) {
	streamPayload, err := os.ReadFile("../../migrations/0063_oidc_security_events.sql")
	if err != nil {
		t.Fatalf("read security event stream migration: %v", err)
	}
	streamSQL := string(streamPayload)
	for _, statement := range []string{
		"CREATE TABLE IF NOT EXISTS gateway_security_event_stream_state",
		"CREATE TABLE IF NOT EXISTS gateway_security_event_verification_challenges",
		"CREATE INDEX IF NOT EXISTS idx_gateway_security_event_challenges_expiry",
	} {
		if !strings.Contains(streamSQL, statement) {
			t.Fatalf("security event stream migration is missing %q", statement)
		}
	}

	connectionPayload, err := os.ReadFile("../../migrations/0064_security_event_connections.sql")
	if err != nil {
		t.Fatalf("read security event connection migration: %v", err)
	}
	connectionSQL := string(connectionPayload)
	for _, statement := range []string{
		"CREATE TABLE IF NOT EXISTS gateway_security_event_connections",
		"management_client_id text",
		"management_credential_ref text",
		"CREATE TABLE IF NOT EXISTS gateway_security_event_connection_idempotency",
	} {
		if !strings.Contains(connectionSQL, statement) {
			t.Fatalf("security event connection migration is missing %q", statement)
		}
	}
}

func TestIdentityConfigurationRevisionMigrationDefinesSecretReferencesAndSingleActiveRevision(t *testing.T) {
	payload, err := os.ReadFile("../../migrations/0065_identity_configuration_revisions.sql")
	if err != nil {
		t.Fatal(err)
	}
	content := string(payload)
	for _, required := range []string{
		"CREATE TABLE IF NOT EXISTS gateway_identity_configuration_revisions",
		"machine_credential_ref text",
		"session_encryption_key_ref text",
		"security_event_configuration_url text",
		"WHERE state = 'active'",
		"CHECK (state IN ('draft','validated','active','superseded','failed'))",
	} {
		if !strings.Contains(content, required) {
			t.Fatalf("identity configuration migration is missing %q", required)
		}
	}
	for _, forbidden := range []string{"machine_secret", "client_secret", "exchange_token text", "onboarding_code text"} {
		if strings.Contains(strings.ToLower(content), forbidden) {
			t.Fatalf("identity configuration migration stores forbidden secret field %q", forbidden)
		}
	}
}

func TestIdentityOnboardingExchangeMigrationStoresOnlySecretReferences(t *testing.T) {
	payload, err := os.ReadFile("../../migrations/0066_identity_onboarding_exchanges.sql")
	if err != nil {
		t.Fatal(err)
	}
	content := strings.ToLower(string(payload))
	for _, required := range []string{
		"create table if not exists gateway_identity_onboarding_exchanges",
		"exchange_token_ref text not null",
		"cleanup_status text not null default 'none'",
		"security_event_credential_handoff_unsafe",
	} {
		if !strings.Contains(content, required) {
			t.Fatalf("identity onboarding migration is missing %q", required)
		}
	}
	for _, forbidden := range []string{"exchange_token text", "onboarding_code", "client_secret", "machine_secret"} {
		if strings.Contains(content, forbidden) {
			t.Fatalf("identity onboarding migration stores forbidden secret field %q", forbidden)
		}
	}
}

func TestIdentityOnboardingMigrationDefinesCancellationLifecycle(t *testing.T) {
	payload, err := os.ReadFile("../../migrations/0066_identity_onboarding_exchanges.sql")
	if err != nil {
		t.Fatal(err)
	}
	content := strings.ToLower(string(payload))
	for _, required := range []string{
		"'cancelled'",
		"cleanup_status text not null default 'none'",
		"cleanup_completed_at",
		"idx_gateway_identity_onboarding_cleanup",
		"last_error_category",
	} {
		if !strings.Contains(content, required) {
			t.Fatalf("identity pairing cancellation migration is missing %q", required)
		}
	}
	for _, forbidden := range []string{"exchange_token text", "client_secret", "machine_secret"} {
		if strings.Contains(content, forbidden) {
			t.Fatalf("identity pairing cancellation migration stores forbidden secret field %q", forbidden)
		}
	}
}

func TestIdentityOnboardingMigrationDefinesFinalErrorCategories(t *testing.T) {
	payload, err := os.ReadFile("../../migrations/0066_identity_onboarding_exchanges.sql")
	if err != nil {
		t.Fatal(err)
	}
	content := strings.ToLower(string(payload))
	for _, required := range []string{
		"gateway_identity_onboarding_error_category_check",
		"security_event_credential_handoff_unsafe",
		"security_event_connection_binding_missing",
		"security_event_connection_binding_unavailable",
		"security_event_connection_binding_invalid",
		"security_event_connection_binding_mismatch",
		"cleanup_security_event_credential_handoff_unsafe",
	} {
		if !strings.Contains(content, required) {
			t.Fatalf("identity pairing error category migration is missing %q", required)
		}
	}
}

func TestIdentityPairingErrorCategoriesExecuteAgainstCurrentSchema(t *testing.T) {
	pool := newIdentityMigrationPostgresTestSchema(t)
	ctx := context.Background()
	for _, migration := range []string{
		"../../migrations/0065_identity_configuration_revisions.sql",
		"../../migrations/0066_identity_onboarding_exchanges.sql",
		"../../migrations/0067_identity_secret_cleanup_queue.sql",
		"../../migrations/0068_identity_pairing_start_reservation.sql",
	} {
		applyIdentityMigrationTestFile(t, ctx, pool, migration)
	}

	revisionID, pairingID := uuid.NewString(), uuid.NewString()
	if _, err := pool.Exec(ctx, `
INSERT INTO gateway_identity_configuration_revisions (
  id,state,auth_center_url,local_tenant_key,public_base_url,web_base_url
) VALUES ($1::uuid,'draft','https://auth.test.example','default',
  'https://gateway.test.example','https://gateway-web.test.example')`, revisionID); err != nil {
		t.Fatalf("seed identity revision: %v", err)
	}
	if _, err := pool.Exec(ctx, `
INSERT INTO gateway_identity_onboarding_exchanges (
  id,revision_id,remote_exchange_id,exchange_token_ref,status,remote_version,expires_at
) VALUES ($1::uuid,$2::uuid,$3::uuid,$4,'credentials_saved',4,now() + interval '30 minutes')`,
		pairingID, revisionID, uuid.NewString(), "identity-exchange-"+pairingID); err != nil {
		t.Fatalf("seed onboarding exchange: %v", err)
	}

	for _, category := range []string{
		"security_event_credential_handoff_unsafe",
		"security_event_connection_binding_missing",
		"security_event_connection_binding_unavailable",
		"security_event_connection_binding_invalid",
		"security_event_connection_binding_mismatch",
		"cleanup_security_event_credential_handoff_unsafe",
	} {
		if _, err := pool.Exec(ctx, `UPDATE gateway_identity_onboarding_exchanges
SET last_error_category=$2 WHERE id=$1::uuid`, pairingID, category); err != nil {
			t.Fatalf("persist supported error category %q: %v", category, err)
		}
	}
	_, err := pool.Exec(ctx, `UPDATE gateway_identity_onboarding_exchanges
SET last_error_category='credential_secret_leaked' WHERE id=$1::uuid`, pairingID)
	requireIdentityMigrationCheckViolation(t, err, "unknown upgraded identity pairing category")
}

func TestIdentitySecretCleanupQueueMigrationStoresOnlyReferences(t *testing.T) {
	payload, err := os.ReadFile("../../migrations/0067_identity_secret_cleanup_queue.sql")
	if err != nil {
		t.Fatal(err)
	}
	content := strings.ToLower(string(payload))
	for _, required := range []string{
		"create table if not exists gateway_identity_secret_cleanup_queue",
		"secret_ref text primary key",
		"not_before timestamptz not null",
		"status text not null default 'pending'",
		"claim_token uuid",
		"lease_expires_at timestamptz",
		"gateway_identity_secret_cleanup_claim_check",
		"idx_gateway_identity_secret_cleanup_due",
	} {
		if !strings.Contains(content, required) {
			t.Fatalf("identity Secret cleanup migration is missing %q", required)
		}
	}
	for _, forbidden := range []string{"secret_value", "client_secret", "machine_secret", "token text"} {
		if strings.Contains(content, forbidden) {
			t.Fatalf("identity Secret cleanup migration stores forbidden value field %q", forbidden)
		}
	}
}

func TestIdentityPairingCancellationLifecycleRejectsPartialStates(t *testing.T) {
	pool := newIdentityMigrationPostgresTestSchema(t)
	ctx := context.Background()
	applyIdentityMigrationTestFile(t, ctx, pool, "../../migrations/0065_identity_configuration_revisions.sql")
	applyIdentityMigrationTestFile(t, ctx, pool, "../../migrations/0066_identity_onboarding_exchanges.sql")

	revisionID, pairingID := uuid.NewString(), uuid.NewString()
	if _, err := pool.Exec(ctx, `
INSERT INTO gateway_identity_configuration_revisions (
  id,state,auth_center_url,local_tenant_key,public_base_url,web_base_url
) VALUES ($1::uuid,'draft','https://auth.test.example','default',
  'https://gateway.test.example','https://gateway-web.test.example')`, revisionID); err != nil {
		t.Fatalf("seed legacy identity revision: %v", err)
	}
	if _, err := pool.Exec(ctx, `
INSERT INTO gateway_identity_onboarding_exchanges (
  id,revision_id,remote_exchange_id,exchange_token_ref,status,remote_version,expires_at,last_error_category
) VALUES ($1::uuid,$2::uuid,$3::uuid,$4,'credentials_saved',4,now() + interval '30 minutes','pairing_step_failed')`,
		pairingID, revisionID, uuid.NewString(), "identity-exchange-"+pairingID); err != nil {
		t.Fatalf("seed legacy onboarding exchange: %v", err)
	}

	var status, cleanupStatus, errorCategory string
	var cancelledAt, cleanupCompletedAt *time.Time
	if err := pool.QueryRow(ctx, `
SELECT status,cleanup_status,last_error_category,cancelled_at,cleanup_completed_at
FROM gateway_identity_onboarding_exchanges WHERE id=$1::uuid`, pairingID).
		Scan(&status, &cleanupStatus, &errorCategory, &cancelledAt, &cleanupCompletedAt); err != nil {
		t.Fatalf("read migrated onboarding exchange: %v", err)
	}
	if status != "credentials_saved" || cleanupStatus != "none" || errorCategory != "pairing_step_failed" ||
		cancelledAt != nil || cleanupCompletedAt != nil {
		t.Fatalf("unexpected migrated exchange status=%q cleanup=%q error=%q cancelled=%v completed=%v",
			status, cleanupStatus, errorCategory, cancelledAt, cleanupCompletedAt)
	}

	_, err := pool.Exec(ctx, `UPDATE gateway_identity_onboarding_exchanges SET status='cancelled' WHERE id=$1::uuid`, pairingID)
	requireIdentityMigrationCheckViolation(t, err, "cancelled status without cleanup intent")

	if _, err := pool.Exec(ctx, `UPDATE gateway_identity_onboarding_exchanges
SET status='cancelled',cleanup_status='pending',cancelled_at=now()
WHERE id=$1::uuid`, pairingID); err != nil {
		t.Fatalf("persist valid pending cleanup lifecycle: %v", err)
	}
	_, err = pool.Exec(ctx, `UPDATE gateway_identity_onboarding_exchanges
SET cleanup_status='completed' WHERE id=$1::uuid`, pairingID)
	requireIdentityMigrationCheckViolation(t, err, "completed cleanup without completion timestamp")

	if _, err := pool.Exec(ctx, `UPDATE gateway_identity_onboarding_exchanges
SET cleanup_status='completed',cleanup_completed_at=now()
WHERE id=$1::uuid`, pairingID); err != nil {
		t.Fatalf("persist valid completed cleanup lifecycle: %v", err)
	}
	_, err = pool.Exec(ctx, `UPDATE gateway_identity_onboarding_exchanges
SET last_error_category='still-invalid' WHERE id=$1::uuid`, pairingID)
	requireIdentityMigrationCheckViolation(t, err, "invalid post-migration error category")

	if err := pool.QueryRow(ctx, `SELECT status,cleanup_status FROM gateway_identity_onboarding_exchanges WHERE id=$1::uuid`, pairingID).
		Scan(&status, &cleanupStatus); err != nil {
		t.Fatalf("read completed cleanup lifecycle: %v", err)
	}
	if status != "cancelled" || cleanupStatus != "completed" {
		t.Fatalf("completed cleanup lifecycle status=%q cleanup=%q", status, cleanupStatus)
	}
}

func applyIdentityMigrationTestFile(t *testing.T, ctx context.Context, pool *pgxpool.Pool, path string) {
	t.Helper()
	payload, err := os.ReadFile(path)
	if err != nil {
		t.Fatalf("read migration %s: %v", path, err)
	}
	tx, err := pool.Begin(ctx)
	if err != nil {
		t.Fatalf("begin migration %s: %v", path, err)
	}
	if _, err := tx.Exec(ctx, string(payload)); err != nil {
		_ = tx.Rollback(ctx)
		t.Fatalf("execute migration %s: %v", path, err)
	}
	if err := tx.Commit(ctx); err != nil {
		t.Fatalf("commit migration %s: %v", path, err)
	}
}

func requireIdentityMigrationCheckViolation(t *testing.T, err error, operation string) {
	t.Helper()
	if err == nil {
		t.Fatalf("%s unexpectedly satisfied migration constraints", operation)
	}
	var postgresError *pgconn.PgError
	if !errors.As(err, &postgresError) || postgresError.Code != "23514" {
		t.Fatalf("%s error=%v, want PostgreSQL check violation", operation, err)
	}
}

func newIdentityMigrationPostgresTestSchema(t *testing.T) *pgxpool.Pool {
	t.Helper()
	databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
	if databaseURL == "" {
		t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run identity migration PostgreSQL integration tests")
	}
	ctx := context.Background()
	admin, err := pgxpool.New(ctx, databaseURL)
	if err != nil {
		t.Fatalf("connect identity migration test database: %v", err)
	}
	var databaseName string
	if err := admin.QueryRow(ctx, `SELECT current_database()`).Scan(&databaseName); err != nil {
		admin.Close()
		t.Fatalf("read identity migration test database name: %v", err)
	}
	if !strings.Contains(strings.ToLower(databaseName), "test") {
		admin.Close()
		t.Fatalf("refusing to use non-test database %q", databaseName)
	}

	schemaName := "gateway_identity_migration_" + strings.ReplaceAll(uuid.NewString(), "-", "")
	schemaIdentifier := pgx.Identifier{schemaName}.Sanitize()
	if _, err := admin.Exec(ctx, `CREATE SCHEMA `+schemaIdentifier); err != nil {
		admin.Close()
		t.Fatalf("create identity migration test schema: %v", err)
	}
	config, err := pgxpool.ParseConfig(databaseURL)
	if err != nil {
		_, _ = admin.Exec(ctx, `DROP SCHEMA IF EXISTS `+schemaIdentifier+` CASCADE`)
		admin.Close()
		t.Fatalf("parse identity migration test database URL: %v", err)
	}
	config.ConnConfig.RuntimeParams["search_path"] = schemaName
	pool, err := pgxpool.NewWithConfig(ctx, config)
	if err != nil {
		_, _ = admin.Exec(ctx, `DROP SCHEMA IF EXISTS `+schemaIdentifier+` CASCADE`)
		admin.Close()
		t.Fatalf("connect identity migration test schema: %v", err)
	}
	t.Cleanup(func() {
		pool.Close()
		if _, err := admin.Exec(context.Background(), `DROP SCHEMA IF EXISTS `+schemaIdentifier+` CASCADE`); err != nil {
			t.Errorf("drop identity migration test schema: %v", err)
		}
		admin.Close()
	})
	return pool
}