MusicServer/internal/testutil/pgtest/pgtest.go

//go:build integration

// Package pgtest starts a PostgreSQL container for integration tests.
//
// Use it from TestMain:
//
//	func TestMain(m *testing.M) { os.Exit(pgtest.Main(m)) }
package pgtest

import (
	"context"
	"fmt"
	"io"
	"log"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"music-server/internal/db"

	"github.com/jackc/pgx/v5/pgxpool"
	"github.com/testcontainers/testcontainers-go"
	"github.com/testcontainers/testcontainers-go/wait"
)

const (
	user     = "postgres"
	password = "postgres"
)

var (
	container testcontainers.Container
	host      string
	port      string
	counter   atomic.Int32
)

// Main starts PostgreSQL, runs the tests and stops PostgreSQL again.
func Main(m *testing.M) int {
	ctx := context.Background()
	c, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
		ContainerRequest: testcontainers.ContainerRequest{
			Image:        "postgres:15-alpine",
			ExposedPorts: []string{"5432/tcp"},
			Env:          map[string]string{"POSTGRES_USER": user, "POSTGRES_PASSWORD": password},
			WaitingFor: wait.ForAll(
				wait.ForLog("database system is ready to accept connections").WithOccurrence(2),
				wait.ForListeningPort("5432/tcp"),
			).WithDeadline(2 * time.Minute),
		},
		Started: true,
	})
	if err != nil {
		log.Printf("failed to start PostgreSQL container (is Docker running?): %v", err)
		return 1
	}
	defer func() {
		if err := c.Terminate(ctx); err != nil {
			log.Printf("failed to stop PostgreSQL container: %v", err)
		}
	}()

	container = c
	host, err = c.Host(ctx)
	if err != nil {
		log.Printf("failed to get container host: %v", err)
		return 1
	}
	mapped, err := c.MappedPort(ctx, "5432/tcp")
	if err != nil {
		log.Printf("failed to get container port: %v", err)
		return 1
	}
	port = mapped.Port()
	return m.Run()
}

// ConnString returns the connection string for a database in the container.
func ConnString(dbname string) string {
	return db.ConnString(host, port, user, password, dbname)
}

// NewDatabase creates an empty database and returns its name. It is dropped when the test ends.
func NewDatabase(t *testing.T) string {
	t.Helper()
	ctx := context.Background()
	name := fmt.Sprintf("test_%d_%d", time.Now().UnixNano()%1_000_000, counter.Add(1))

	admin, err := pgxpool.New(ctx, ConnString("postgres"))
	if err != nil {
		t.Fatalf("connect to postgres: %v", err)
	}
	defer admin.Close()
	if _, err := admin.Exec(ctx, "CREATE DATABASE "+name); err != nil {
		t.Fatalf("create database: %v", err)
	}
	t.Cleanup(func() {
		admin, err := pgxpool.New(ctx, ConnString("postgres"))
		if err != nil {
			return
		}
		defer admin.Close()
		admin.Exec(ctx, "DROP DATABASE IF EXISTS "+name+" WITH (FORCE)")
	})
	return name
}

// NewMigratedDatabase creates a database with all migrations applied.
func NewMigratedDatabase(t *testing.T) *db.Database {
	t.Helper()
	database, err := db.NewDatabase(context.Background(), ConnString(NewDatabase(t)))
	if err != nil {
		t.Fatalf("open database: %v", err)
	}
	t.Cleanup(database.Close)
	if err := database.RunMigrations(); err != nil {
		t.Fatalf("run migrations: %v", err)
	}
	return database
}

// LoadSQLFile runs a SQL file with psql inside the container, as the superuser.
func LoadSQLFile(t *testing.T, dbname, path string) {
	t.Helper()
	ctx := context.Background()
	target := "/tmp/" + dbname + ".sql"
	if err := container.CopyFileToContainer(ctx, path, target, 0o644); err != nil {
		t.Fatalf("copy %s to container: %v", path, err)
	}
	code, reader, err := container.Exec(ctx, []string{"psql", "-U", user, "-d", dbname, "-v", "ON_ERROR_STOP=1", "-q", "-f", target})
	if err != nil {
		t.Fatalf("run psql: %v", err)
	}
	output, _ := io.ReadAll(reader)
	if code != 0 {
		t.Fatalf("psql exited with %d: %s", code, output)
	}
}

// CreateRole creates a login role, used by dumps that reference their original owner.
func CreateRole(t *testing.T, name string) {
	t.Helper()
	ctx := context.Background()
	admin, err := pgxpool.New(ctx, ConnString("postgres"))
	if err != nil {
		t.Fatalf("connect to postgres: %v", err)
	}
	defer admin.Close()
	_, err = admin.Exec(ctx, "CREATE ROLE "+name+" LOGIN")
	if err != nil && !strings.Contains(err.Error(), "already exists") {
		t.Fatalf("create role: %v", err)
	}
}