436 lines
12 KiB
Go
436 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"
|
|
)
|
|
|
|
func dbHost() string {
|
|
if h := os.Getenv("TEST_DB_HOST"); h != "" {
|
|
return h
|
|
}
|
|
if h := os.Getenv("POSTGRES_HOST"); h != "" {
|
|
return h
|
|
}
|
|
return "localhost"
|
|
}
|
|
|
|
// 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()
|
|
|
|
adminDSN := fmt.Sprintf("postgres://myuser:mypassword@%s:5432/mydb?sslmode=disable", dbHost())
|
|
adminPool, err := pgxpool.New(ctx, adminDSN)
|
|
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@%s:5432/%s?sslmode=disable", dbHost(), 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()
|
|
adminDSN := fmt.Sprintf("postgres://myuser:mypassword@%s:5432/mydb?sslmode=disable", dbHost())
|
|
adminPool, err := pgxpool.New(ctx, adminDSN)
|
|
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 { // #nosec G304
|
|
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")
|
|
}
|