Crussell/backend/testutils/testdb/testdb.go

//go:build test
// +build test

package testdb

import (
	"context"
	"fmt"
	"os"
	"path/filepath"
	"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()

	// Check if database already has tables or types
	var err error
	var typeCount int
	err = pool.QueryRow(ctx, "SELECT COUNT(*) FROM pg_type WHERE typnamespace = (SELECT oid FROM pg_namespace WHERE nspname = 'public')").Scan(&typeCount)
	if err == nil && typeCount > 0 {
		t.Log("Database already has types, dropping and recreating for clean state...")
		
		// 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",
		}
		for _, stmt := range typeDrops {
			if _, err := pool.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{
			"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",
			"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 := pool.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 := pool.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")
	}

	// Simple migration: just create tables that don't exist
	// Note: This doesn't handle stored procedures properly, but the database
	// should already be set up with the correct schema
	t.Log("Running migration...")

	// Execute the schema SQL
	_, err = pool.Exec(ctx, schemaSQL)
	if err != nil {
		t.Fatalf("Failed to execute migration: %v", err)
	}
}

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

    tables := []string{
        "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",
        "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 := pool.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")
}