//go:build test // +build test package testdb import ( "context" "fmt" "os" "path/filepath" "strings" "testing" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgxpool" ) const defaultTestDSN = "postgres://myuser:mypassword@localhost:5432/crussell_test?sslmode=disable" func Pool(t *testing.T) *pgxpool.Pool { t.Helper() dsn := os.Getenv("TEST_DB_DSN") if dsn == "" { dsn = defaultTestDSN } ctx := context.Background() pool, err := pgxpool.New(ctx, dsn) if err != nil { t.Fatalf("Failed to connect to test database: %v", err) } if err := pool.Ping(ctx); err != nil { t.Fatalf("Failed to ping test database: %v", err) } return pool } func NewPool(dsn string) (*pgxpool.Pool, error) { if dsn == "" { dsn = defaultTestDSN } ctx := context.Background() pool, err := pgxpool.New(ctx, dsn) if err != nil { return nil, fmt.Errorf("failed to create pool: %w", err) } if err := pool.Ping(ctx); err != nil { return nil, fmt.Errorf("failed to ping database: %w", err) } return pool, nil } func Migrate(t *testing.T, pool *pgxpool.Pool) { t.Helper() ctx := context.Background() // Acquire a dedicated connection so advisory lock, migration, and unlock all use the same session conn, err := pool.Acquire(ctx) if err != nil { t.Fatalf("Failed to acquire connection for migration: %v", err) } defer conn.Release() _, err = conn.Exec(ctx, "SELECT pg_advisory_lock(1337)") if err != nil { t.Fatalf("Failed to acquire migration advisory lock: %v", err) } defer conn.Exec(ctx, "SELECT pg_advisory_unlock(1337)") // Check if database already has tables or types var typeCount int err = conn.QueryRow(ctx, "SELECT COUNT(*) FROM pg_type WHERE typnamespace = (SELECT oid FROM pg_namespace WHERE nspname = 'public')").Scan(&typeCount) if err == nil && typeCount > 0 { // Drop TYPES FIRST (they have CASCADE dependencies on tables) typeDrops := []string{ "DROP TYPE IF EXISTS account_role CASCADE", "DROP TYPE IF EXISTS account_type CASCADE", "DROP TYPE IF EXISTS payment_type CASCADE", "DROP TYPE IF EXISTS payment_method CASCADE", "DROP TYPE IF EXISTS payment_status CASCADE", "DROP TYPE IF EXISTS booking_status CASCADE", "DROP TYPE IF EXISTS verification_purpose CASCADE", "DROP TYPE IF EXISTS admin_notification_reason CASCADE", "DROP TYPE IF EXISTS campaign_type CASCADE", "DROP TYPE IF EXISTS milestone_type CASCADE", "DROP TYPE IF EXISTS milestone_unit CASCADE", "DROP TYPE IF EXISTS discount_campaign_scope CASCADE", "DROP TYPE IF EXISTS discount_campaign_status CASCADE", "DROP TYPE IF EXISTS till_item_type CASCADE", } for _, stmt := range typeDrops { if _, err := conn.Exec(ctx, stmt); err != nil { t.Logf("Warning dropping type: %v (expected if using IF EXISTS)", err) } } // Drop all tables, sequences, and views in correct order dropOrder := []string{ "till_sales", "admin_audit_log", "gift_card_transactions", "gift_card_expired_balances", "booking_discounts", "loyalty_redemptions", "discount_campaigns", "forgiven_no_shows", "admin_notifications", "user_notification_preferences", "user_referrals", "booking_custom_services", "booking_services", "refunds", "payments", "user_saved_cards", "square_deposits", "affiliate_payouts", "financial_aggregates", "gift_cards", "user_giftcard_balances", "bookings", "user_patch_tests", "patch_tests", "booking_edit_requests", "services", "custom_services", "verification_codes", "user_social_logins", "users", "images", "tags", "time_blockers", "working_hours", "exceptional_group_applications", "exceptional_working_hours", "exceptional_working_hours_groups", "business_settings", } for _, table := range dropOrder { stmt := fmt.Sprintf("DROP TABLE IF EXISTS %s CASCADE", table) if _, err := conn.Exec(ctx, stmt); err != nil { t.Logf("Warning dropping table %s: %v (expected if using IF EXISTS)", table, err) } } // Drop sequences last seqDrops := []string{ "DROP SEQUENCE IF EXISTS invoice_number_seq", "DROP SEQUENCE IF EXISTS tags_id_seq", "DROP SEQUENCE IF EXISTS exceptional_working_hours_id_seq", "DROP SEQUENCE IF EXISTS exceptional_working_hours_groups_id_seq", } for _, stmt := range seqDrops { if _, err := conn.Exec(ctx, stmt); err != nil { t.Logf("Warning dropping sequence: %v", err) } } } paths := []string{ "../../../init-scripts/init-script.sql", "../../init-scripts/init-script.sql", "../init-scripts/init-script.sql", "init-scripts/init-script.sql", } var schemaSQL string for _, p := range paths { if data, err := os.ReadFile(p); err == nil { schemaSQL = string(data) break } } if schemaSQL == "" { t.Fatal("Could not find init-script.sql in any expected location") } // Split and execute the schema SQL statement-by-statement for robust execution and diagnostic visibility statements := splitSQLStatements(schemaSQL) for i, stmt := range statements { stmt = strings.TrimSpace(stmt) if stmt == "" { continue } firstLine := strings.Split(stmt, "\n")[0] if len(firstLine) > 80 { firstLine = firstLine[:80] + "..." } _, err = conn.Exec(ctx, stmt) if err != nil { t.Fatalf("Failed to execute migration statement [%d]: %s\nError: %v", i+1, firstLine, err) } } } func splitSQLStatements(sql string) []string { var statements []string var currentStatement strings.Builder inDollarBlock := false lines := strings.Split(sql, "\n") for _, line := range lines { trimmed := strings.TrimSpace(line) if trimmed == "" { continue } if strings.HasPrefix(trimmed, "--") { continue } currentStatement.WriteString(line) currentStatement.WriteString("\n") // Count occurrences of $$ in the line (handles single-line $$ blocks too) count := strings.Count(line, "$$") if count%2 == 1 { inDollarBlock = !inDollarBlock } if !inDollarBlock && strings.HasSuffix(trimmed, ";") { statements = append(statements, currentStatement.String()) currentStatement.Reset() } } if currentStatement.Len() > 0 { statements = append(statements, currentStatement.String()) } return statements } func Tx(t *testing.T, pool *pgxpool.Pool) pgx.Tx { t.Helper() ctx := context.Background() tx, err := pool.Begin(ctx) if err != nil { t.Fatalf("Failed to begin transaction: %v", err) } return tx } func TxWithRollback(t *testing.T, pool *pgxpool.Pool) (pgx.Tx, func()) { tx := Tx(t, pool) return tx, func() { tx.Rollback(context.Background()) } } func TruncateTables(t *testing.T, pool *pgxpool.Pool) { t.Helper() ctx := context.Background() // Acquire a dedicated connection so advisory lock, truncations, and unlock all use the same session conn, err := pool.Acquire(ctx) if err != nil { t.Fatalf("Failed to acquire connection for truncation: %v", err) } defer conn.Release() _, err = conn.Exec(ctx, "SELECT pg_advisory_lock(1338)") if err != nil { t.Fatalf("Failed to acquire truncate advisory lock: %v", err) } defer conn.Exec(ctx, "SELECT pg_advisory_unlock(1338)") tables := []string{ "till_sales", "admin_audit_log", "gift_card_transactions", "gift_card_expired_balances", "booking_discounts", "loyalty_redemptions", "discount_campaigns", "forgiven_no_shows", "user_social_logins", "verification_codes", "booking_custom_services", "booking_services", "refunds", "payments", "user_saved_cards", "square_deposits", "affiliate_payouts", "financial_aggregates", "gift_cards", "user_giftcard_balances", "bookings", "booking_edit_requests", "user_patch_tests", "patch_tests", "services", "custom_services", "admin_notifications", "user_referrals", "user_notification_preferences", "time_blockers", "working_hours", "exceptional_group_applications", "exceptional_working_hours", "exceptional_working_hours_groups", "business_settings", "users", "images", "tags", } for _, table := range tables { _, err := conn.Exec(ctx, fmt.Sprintf("TRUNCATE TABLE %s CASCADE", table)) if err != nil { t.Logf("Warning: could not truncate %s: %v", table, err) } } } func FindInitScript() (string, error) { cwd, err := os.Getwd() if err != nil { cwd = "" } paths := []string{ "../../../init-scripts/init-script.sql", "../../init-scripts/init-script.sql", "../init-scripts/init-script.sql", "init-scripts/init-script.sql", } if cwd != "" { paths = append(paths, filepath.Join(cwd, "..", "..", "init-scripts", "init-script.sql")) } for _, p := range paths { if _, err := os.Stat(p); err == nil { return p, nil } } return "", fmt.Errorf("could not find init-script.sql") }