//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" 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() // Check if database already has tables by checking for the users table ctx := context.Background() var tableCount int err := pool.QueryRow(ctx, "SELECT COUNT(*) FROM pg_tables WHERE tablename = 'users'").Scan(&tableCount) if err == nil && tableCount > 0 { // Tables already exist, skip migration t.Log("Database already has tables, skipping migration") return } 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...") } 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{ "user_social_logins", "verification_codes", "booking_services", "payments", "bookings", "user_service_patch_tests", "services", "admin_notifications", "user_referrals", "user_notification_preferences", "working_hours", "exceptional_working_hours", "exceptional_working_hours_groups", "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") }