Files
Crussell/backend/testutils/testdb/testdb.go
T
popertotsandSisyphus e4b9003439 refactor(handlers): migrate remaining backend handlers to clock.Now() and transaction patterns
Apply clock.Now() migration, transaction wrapping, and minor refactors across admin, scheduling, today, user, auth handler, notifications, webhooks, services, portfolio, ratelimit, testutils, and main.go.

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
2026-06-24 23:43:50 +01:00

424 lines
12 KiB
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/pgconn"
"github.com/jackc/pgx/v5/pgxpool"
)
// Querier is a minimal interface satisfied by *db.PoolProxy, *pgxpool.Pool,
// and pgx.Tx. Defined locally to avoid an import cycle (db → testdb → db).
type Querier interface {
Exec(ctx context.Context, sql string, args ...any) (pgconn.CommandTag, error)
}
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)
poolCfg, err := pgxpool.ParseConfig(testDSN)
if err != nil {
log.Fatalf("testdb: failed to parse config: %v", err)
}
poolCfg.MaxConns = 16
poolCfg.ConnConfig.DefaultQueryExecMode = pgx.QueryExecModeSimpleProtocol
poolCfg.ConnConfig.RuntimeParams["timezone"] = "UTC"
pool, err := pgxpool.NewWithConfig(ctx, poolCfg)
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()
poolCfg, err := pgxpool.ParseConfig(dsn)
if err != nil {
t.Fatalf("Failed to parse config: %v", err)
}
poolCfg.ConnConfig.RuntimeParams["timezone"] = "UTC"
pool, err := pgxpool.NewWithConfig(ctx, poolCfg)
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()
poolCfg, err := pgxpool.ParseConfig(dsn)
if err != nil {
return nil, fmt.Errorf("failed to parse config: %w", err)
}
poolCfg.ConnConfig.RuntimeParams["timezone"] = "UTC"
pool, err := pgxpool.NewWithConfig(ctx, poolCfg)
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
}
// SeedBaseline populates the test database with reference data that is needed
// by most tests but never changes across test runs (working hours, business
// settings, etc.). Called once per TestMain AFTER migration but BEFORE m.Run().
// Data is committed at the pool level and visible inside all per-test transactions
// (PostgreSQL Read Committed isolation).
func SeedBaseline(pool *pgxpool.Pool) {
ctx := context.Background()
// Working hours: Mon-Sun 08:00-20:00, all open.
// Used by bookings, admin, and most handler tests that schedule appointments.
hours := []struct {
weekday int
startTime string
endTime string
isOpen bool
}{
{0, "08:00", "20:00", true},
{1, "08:00", "20:00", true},
{2, "08:00", "20:00", true},
{3, "08:00", "20:00", true},
{4, "08:00", "20:00", true},
{5, "08:00", "20:00", true},
{6, "08:00", "20:00", true},
}
for _, h := range hours {
_, err := pool.Exec(ctx, `
INSERT INTO working_hours (weekday, start_time, end_time, is_open)
VALUES ($1, $2, $3, $4)
ON CONFLICT (weekday) DO UPDATE SET start_time = $2, end_time = $3, is_open = $4
`, h.weekday, h.startTime, h.endTime, h.isOpen)
if err != nil {
log.Fatalf("testdb: failed to seed working hours: %v", err)
}
}
// Business settings: required by admin/settings tests and booking deposit logic.
_, err := pool.Exec(ctx, `
INSERT INTO business_settings (business_name, business_address, currency_code, gift_card_expiry_months, voucher_type)
VALUES ('Test Salon', '123 Test St', 'GBP', 12, 'SPV')
ON CONFLICT DO NOTHING
`)
if err != nil {
log.Fatalf("testdb: failed to seed business settings: %v", err)
}
}
// SeedBaselineScheduling is like SeedBaseline but uses working hours appropriate
// for the scheduling package's tests (Mon-Fri 09:00-17:00, Sat 10:00-16:00, Sun closed).
func SeedBaselineScheduling(pool *pgxpool.Pool) {
ctx := context.Background()
hours := []struct {
weekday int
startTime string
endTime string
isOpen bool
}{
{0, "09:00", "17:00", true}, // Monday
{1, "09:00", "17:00", true}, // Tuesday
{2, "09:00", "17:00", true}, // Wednesday
{3, "09:00", "17:00", true}, // Thursday
{4, "09:00", "17:00", true}, // Friday
{5, "10:00", "16:00", true}, // Saturday
{6, "09:00", "17:00", false}, // Sunday (closed)
}
for _, h := range hours {
_, err := pool.Exec(ctx, `
INSERT INTO working_hours (weekday, start_time, end_time, is_open)
VALUES ($1, $2, $3, $4)
ON CONFLICT (weekday) DO UPDATE SET start_time = $2, end_time = $3, is_open = $4
`, h.weekday, h.startTime, h.endTime, h.isOpen)
if err != nil {
log.Fatalf("testdb: failed to seed scheduling working hours: %v", err)
}
}
// Business settings same as baseline.
_, err := pool.Exec(ctx, `
INSERT INTO business_settings (business_name, business_address, currency_code, gift_card_expiry_months, voucher_type)
VALUES ('Test Salon', '123 Test St', 'GBP', 12, 'SPV')
ON CONFLICT DO NOTHING
`)
if err != nil {
log.Fatalf("testdb: failed to seed business settings: %v", err)
}
}
// 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 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")
}