Crussell/backend/testutils/testdb/testdb.go

//go:build test
// +build test

package testdb

import (
	"context"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

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

const (
	defaultTestDSN   = "postgres://myuser:mypassword@localhost:5432/crussell_test?sslmode=disable"
	adminDatabaseDSN = "postgres://myuser:mypassword@localhost:5432/mydb?sslmode=disable"
)

// CreateTestDatabase creates a fresh isolated test database, runs the schema migration,
// and returns a pool connected to it. If a database with the same name exists, it's
// dropped first. This enables parallel test execution across packages since each
// package gets its own database running migrations in parallel.
func CreateTestDatabase(dbName string) *pgxpool.Pool {
	ctx := context.Background()

	adminPool, err := pgxpool.New(ctx, adminDatabaseDSN)
	if err != nil {
		log.Fatalf("testdb: failed to connect to admin database: %v", err)
	}
	defer adminPool.Close()

	safeName := pgx.Identifier{dbName}.Sanitize()

	// Drop existing database with retries (terminate connections first)
	for attempt := 0; attempt < 3; attempt++ {
		_, _ = adminPool.Exec(ctx, fmt.Sprintf(`
			SELECT pg_terminate_backend(pid)
			FROM pg_stat_activity
			WHERE datname = '%s' AND pid != pg_backend_pid()
		`, dbName))

		_, err = adminPool.Exec(ctx, fmt.Sprintf("DROP DATABASE IF EXISTS %s", safeName))
		if err == nil {
			break
		}
		log.Printf("testdb: retry %d dropping %s: %v", attempt+1, dbName, err)
	}
	if err != nil {
		log.Fatalf("testdb: failed to drop database %s after retries: %v", dbName, err)
	}

	_, err = adminPool.Exec(ctx, fmt.Sprintf("CREATE DATABASE %s", safeName))
	if err != nil {
		log.Fatalf("testdb: failed to create database %s: %v", dbName, err)
	}
	adminPool.Close()

	testDSN := fmt.Sprintf("postgres://myuser:mypassword@localhost:5432/%s?sslmode=disable", dbName)
	pool, err := pgxpool.New(ctx, testDSN)
	if err != nil {
		log.Fatalf("testdb: failed to connect to %s: %v", dbName, err)
	}

	if err := pool.Ping(ctx); err != nil {
		log.Fatalf("testdb: failed to ping %s: %v", dbName, err)
	}

	if err := migratePool(pool); err != nil {
		log.Fatalf("testdb: migration failed: %v", err)
	}

	return pool
}

// DestroyTestDatabase closes the pool (if non-nil) and drops the test database.
// Retries aggressively with exponential backoff, pg_terminate_backend between
// attempts, and a final DROP DATABASE WITH (FORCE) on PostgreSQL 13+.
func DestroyTestDatabase(pool *pgxpool.Pool, dbName string) {
	if dbName == "" {
		return
	}

	if pool != nil {
		pool.Close()
	}

	ctx := context.Background()
	adminPool, err := pgxpool.New(ctx, adminDatabaseDSN)
	if err != nil {
		log.Printf("testdb: warning: failed to connect to drop %s: %v", dbName, err)
		return
	}
	defer adminPool.Close()

	safeName := pgx.Identifier{dbName}.Sanitize()

	// Aggressive retry with pg_terminate_backend between attempts and final FORCE option.
	for attempt := 0; attempt < 10; attempt++ {
		// Kill all connections to the target database.
		_, _ = adminPool.Exec(ctx, fmt.Sprintf(`
			SELECT pg_terminate_backend(pid)
			FROM pg_stat_activity
			WHERE datname = '%s' AND pid != pg_backend_pid()
		`, dbName))

		// Try DROP DATABASE WITH (FORCE) first (PG13+) — this kills remaining connections atomically.
		_, err = adminPool.Exec(ctx, fmt.Sprintf("DROP DATABASE IF EXISTS %s WITH (FORCE)", safeName))
		if err == nil {
			return
		}

		// Fallback to plain DROP (may work if pg_terminate_backend was sufficient).
		_, err = adminPool.Exec(ctx, fmt.Sprintf("DROP DATABASE IF EXISTS %s", safeName))
		if err == nil {
			return
		}

		// Exponential backoff: 100ms, 200ms, 400ms, 800ms, 1.6s, 3.2s, 6.4s, 12.8s, 25.6s, 51.2s
		wait := time.Duration(100*(1<<attempt)) * time.Millisecond
		log.Printf("testdb: retry %d dropping %s: %v (retrying in %v)", attempt+1, dbName, err, wait)
		time.Sleep(wait)
	}

	log.Printf("testdb: warning: failed to drop database %s after 10 attempts: %v", dbName, err)
}

// Pool creates a new test database pool using TEST_DB_DSN env var or default DSN.
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
}

// migratePool executes the full schema migration on the given pool.
// This internal version returns errors directly for use by CreateTestDatabase.
func migratePool(pool *pgxpool.Pool) error {
	ctx := context.Background()
	conn, err := pool.Acquire(ctx)
	if err != nil {
		return fmt.Errorf("failed to acquire connection for migration: %w", err)
	}
	defer conn.Release()
	return migrateSchema(ctx, conn)
}

// migrateSchema executes the schema migration on an existing connection.
// Returns an error if any step fails.
// NOTE: This is always called on a fresh database (created via CreateTestDatabase),
// so no defensive cleanup of types/tables/sequences is needed.
func migrateSchema(ctx context.Context, conn *pgxpool.Conn) error {
	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 == "" {
		return fmt.Errorf("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 {
			return fmt.Errorf("failed to execute migration statement [%d]: %s\nError: %w", i+1, firstLine, err)
		}
	}
	return nil
}

func Migrate(t *testing.T, pool *pgxpool.Pool) {
	t.Helper()
	if err := migratePool(pool); err != nil {
		t.Fatal(err.Error())
	}
}

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()

	conn, err := pool.Acquire(ctx)
	if err != nil {
		t.Fatalf("Failed to acquire connection for truncation: %v", err)
	}
	defer conn.Release()

	// Single TRUNCATE with all tables — one round-trip instead of ~44.
	// CASCADE handles FK dependencies so order doesn't matter.
	// Advisory lock removed: each package has its own database now.
	_, err = conn.Exec(ctx, `
		TRUNCATE TABLE
			till_sales, admin_audit_log, gift_card_transactions,
			gift_card_expired_balances, booking_discounts, loyalty_redemptions,
			discount_campaigns, forgiven_no_shows, name_history, referral_discounts,
			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,
			login_audit, refresh_tokens, revoked_jtis
		CASCADE
	`)
	if err != nil {
		t.Logf("Warning: truncation failed: %v", 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")
}