pulley/internal/testutil/db.go

package testutil

import (
	"context"
	"database/sql"
	"fmt"
	"net/url"
	"os"
	"testing"
	"time"

	"git.runcible.io/learning/pulley/internal/config"
	"git.runcible.io/learning/pulley/internal/data"
	"git.runcible.io/learning/pulley/internal/database"
	"git.runcible.io/learning/pulley/migrations"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

// SetupTestDB provides a connection pool to a freshly created temporary database.
//
// Environment Variables:
//   - PULLEY_INTEGRATION_DATABASE_URI: The connection string for the admin/maintenance DB.
//     Defaults to: postgres://pulley:pulley@localhost:5432/postgres?sslmode=disable
//   - PULLEY_ALLOW_REMOTE_INTEGRATION_TEST: Must be 'true' if the URI host is not 'localhost' or '127.0.0.1'.
//
// Safety:
//
//	If the resolved URI points to a remote host and the safety flag is not enabled, the test
//	will fail immediately to prevent accidental execution against production or staging environments.
//
// Lifecycle:
//
//	The function uses t.Cleanup to automatically close the pool and drop the temporary
//	database, ensuring no resource leaks between test runs.
func SetupTestDB(t *testing.T) *pgxpool.Pool {
	t.Helper()

	ctxTimeout, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	intDBUri, ok := os.LookupEnv("PULLEY_INTEGRATION_DATABASE_URI")
	if !ok {
		intDBUri = "postgres://pulley:pulley@localhost:5434/postgres?sslmode=disable"
		t.Log("PULLEY_INTEGRATION_DATABASE_URI not set. Using localhost default")
	}

	// check for non-local hostname
	parsed, err := url.Parse(intDBUri)
	if err != nil {
		t.Fatal(err)
	}

	if parsed.Hostname() != "localhost" && parsed.Hostname() != "127.0.0.1" {
		if os.Getenv("PULLEY_ALLOW_REMOTE_INTEGRATION_TEST") != "true" {
			t.Fatalf("Attempting to run integration tests against a non-local DB (%s). "+
				"Set PULLEY_ALLOW_REMOTE_INTEGRATION_TEST=true to bypass.", parsed.Hostname())
		}
	}

	// verify connection
	dbAdmin, err := sql.Open("pgx", intDBUri)
	if err != nil {
		t.Log("Hint: check if `make start-local` is running a postgres container")
		t.Fatalf("Failed to connect to database: %s", err.Error())
	}
	t.Cleanup(func() {
		dbAdmin.Close()
	})

	err = dbAdmin.PingContext(ctxTimeout)
	if err != nil {
		t.Log("Hint: check if `make start-local` is running a postgres container")
		t.Fatalf("Failed to connect to database: %s", err)
	}

	// create and migrate new db
	testDbName := fmt.Sprintf("test_db_%d", time.Now().UnixNano())
	t.Logf("Creating integration test database: %s", testDbName)
	query := fmt.Sprintf("CREATE DATABASE %s", testDbName)

	_, err = dbAdmin.ExecContext(ctxTimeout, query)
	if err != nil {
		t.Fatalf("Created to create database %s: %s", testDbName, err)
	}

	testCreds := *parsed
	testCreds.Path = fmt.Sprintf("/%s", testDbName)

	// migrations
	migrateDb, err := sql.Open("pgx", testCreds.String())
	if err != nil {
		t.Fatalf("Failed to open  migrations connection to %s: %s", testDbName, err)
	}
	defer migrateDb.Close()

	err = migrateDb.PingContext(ctxTimeout)
	if err != nil {
		t.Fatalf("Failed to connect to database: %s", err)
	}

	t.Logf("Migrating database %s", testDbName)
	err = migrations.Migrate(migrateDb)
	if err != nil {
		t.Fatalf("Failed to migrate %s: %s", testDbName, err)
	}
	migrateDb.Close()

	testPool, err := database.OpenPgPool(ctxTimeout, config.ServiceConfig{DatabaseUri: testCreds.String()})

	t.Cleanup(func() {
		t.Logf("Ensuring test pool is closed")
		testPool.Close()

		t.Logf("Cleaning up database %s", testDbName)

		if t.Failed() && os.Getenv("PULLEY_PRESERVE_FAILED_DB") == "true" {
			t.Logf("Test failed; skipping db cleanup for %s", testDbName)
			return
		}

		ctxCleanup, cleanupCancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cleanupCancel()

		query := fmt.Sprintf("DROP DATABASE IF EXISTS %s WITH (FORCE)", testDbName)

		_, err = dbAdmin.ExecContext(ctxCleanup, query)
		if err != nil {
			t.Errorf("Error in dropping database %s", testDbName)
		}
	})

	return testPool
}

// SeedMovie inserts movies into the movies database table
func SeedMovies(t *testing.T, pool *pgxpool.Pool, movies []data.Movie) {
	t.Helper()

	ctxTimeout, cancel := context.WithTimeout(context.Background(), 5*time.Second)

	defer cancel()

	tx, err := pool.BeginTx(ctxTimeout, pgx.TxOptions{})
	if err != nil {
		t.Fatalf("Failed to begin seed transaction: %v", err)
	}

	defer tx.Rollback(ctxTimeout)

	query := `
        INSERT INTO movies (id, title, year, runtime, genres, version) 
        VALUES ($1, $2, $3, $4, $5, $6)`

	for _, m := range movies {
		_, err := tx.Exec(ctxTimeout, query, m.ID, m.Title, m.Year, m.Runtime, m.Genres, m.Version)
		if err != nil {
			t.Fatalf("failed to insert movie %q: %v", m.Title, err)
		}
	}

	if err = tx.Commit(ctxTimeout); err != nil {
		t.Fatalf("Failed to commit seed transaction: %v", err)
	}

}