Crussell/backend/testutils/testdb/testdb.go

//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",
			"gift_card_transactions",
			"gift_card_expired_balances",
			"booking_discounts",
			"loyalty_redemptions",
			"discount_campaigns",
			"forgiven_no_shows",
			"admin_notifications",
			"user_notification_preferences",
			"user_referrals",
			"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",
			"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",
		"gift_card_transactions",
		"gift_card_expired_balances",
		"booking_discounts",
		"loyalty_redemptions",
		"discount_campaigns",
		"forgiven_no_shows",
		"user_social_logins",
		"verification_codes",
		"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",
		"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")
}