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/database" "git.runcible.io/learning/pulley/migrations" "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() // TODO envvar check could be beneficial if DB needs to be inspected // because of failure t.Logf("Cleaning up database %s", testDbName) 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 }