Postgres test isolation helper

postgres_test.go

package pgtest

import (
	"context"
	"fmt"
	"strings"
	"sync"
	"testing"

	"github.com/docker/go-connections/nat"
	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
	_ "github.com/jackc/pgx/v5/stdlib"
	"github.com/stretchr/testify/require"
	"github.com/testcontainers/testcontainers-go"
	"github.com/testcontainers/testcontainers-go/wait"
)

const (
	testPassword = "hunter2"
	testUser     = "postgres"
	// This is within the container, not what is bound on the host
	testPort = "5432/tcp"
)

// To reduce impact to test startup times, run at most one postgres container.
// Each test should create a unique SCHEMA and use it as the search path for
// connections.
var startPostgres = sync.OnceValues(func() (testcontainers.Container, error) {
	ctx := context.Background()
	return testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
		ContainerRequest: testcontainers.ContainerRequest{
			Image: "postgres:16",
			Env: map[string]string{
				"POSTGRES_PASSWORD": testPassword,
				"POSTGRES_USER":     testUser,
			},
			ExposedPorts: []string{testPort},
			WaitingFor:   wait.ForSQL(testPort, "pgx", getURL),
		},
		Started: true,
	})
})

// New returns a postgres pool that's connected to a live test instance, inside a SCHEMA
// created just for the given test.
//
// Resources are cleaned up on exit.
func New(t testing.TB) *pgxpool.Pool {
	container, err := startPostgres()
	require.NoError(t, err)

	ctx := context.Background()
	addr, err := container.Endpoint(ctx, "")
	require.NoError(t, err)

	parts := strings.Split(addr, ":")
	require.Len(t, parts, 2)

	// DSN to initialize the connection for setup
	dsn := getURL(parts[0], nat.Port(parts[1]+"/tcp"))

	// Connect to DB, create schema for test, then modify DSN to include the
	// new schema in the search_path
	schema := uniqueSchema()
	createSchema(t, ctx, dsn, schema)
	dsn += " search_path=" + schema

	pcfg, err := pgxpool.ParseConfig(dsn)
	require.NoError(t, err)

	pool, err := pgxpool.NewWithConfig(ctx, pcfg)
	require.NoError(t, err)

	t.Cleanup(func() {
		pool.Close()
	})
	return pool
}

// createSchema by connecting to the database identified by 'dsn' and clean it up when
// the given test exits
func createSchema(t testing.TB, ctx context.Context, dsn string, name string) {
	cfg, err := pgx.ParseConfig(dsn)
	require.NoError(t, err)

	conn, err := pgx.ConnectConfig(ctx, cfg)
	require.NoError(t, err)

	_, err = conn.Exec(ctx, "CREATE SCHEMA IF NOT EXISTS "+name+";")
	require.NoError(t, err)

	t.Cleanup(func() {
		_, err = conn.Exec(ctx, "DROP SCHEMA "+name+" CASCADE;")
		require.NoError(t, err)
		err = conn.Close(ctx)
		require.NoError(t, err)
	})
}

// uniqueSchema returns a valid Postgres identifier that can be used to isolate
// tests in their own schema
func uniqueSchema() string {
	// Valid, unquoted identifiers must start with a letter and not have hyphens
	return "t_" + strings.ReplaceAll(uuid.NewString(), "-", "_")
}

// getURL returns a DSN to connect to the test container, given it's host and port
func getURL(host string, port nat.Port) string {
	return fmt.Sprintf(
		"user=%s password=%s host=%s port=%d dbname=%s",
		testUser,
		testPassword,
		host,
		port.Int(),
		testUser,
	)
}

Honestly, the uniqueSchema() function could be replaced with the test's fully-qualified name instead.