query-orchestration/internal/test/database.go

package test

import (
	"context"
	"fmt"
	"queryorchestration/internal/database"
	"testing"

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

type ExternalDatabase struct {
	Name     string
	Host     string
	Port     int
	User     string
	Password string
}

type Database struct {
	Pool      *pgxpool.Pool
	Container testcontainers.Container
	External  *ExternalDatabase
}

type CreateDatabaseConfig struct {
	Network    *testcontainers.DockerNetwork
	Migrations *database.MigrationConfig
	NoPool     bool
}

func CreateDB(t *testing.T, ctx context.Context, cfg *CreateDatabaseConfig) (*Database, func()) {
	alias := "postgres"

	port, err := nat.NewPort("tcp", "5432")
	if err != nil {
		t.Fatalf("Failed to create port: %v", err)
	}

	name := "queryorchestration"
	pass := "pass"
	user := "postgres"

	req := testcontainers.ContainerRequest{
		Image: "postgres:17.2-alpine3.21",
		Env: map[string]string{
			"POSTGRES_DB":       name,
			"POSTGRES_USER":     user,
			"POSTGRES_PASSWORD": pass,
		},
		ExposedPorts: []string{port.Port()},
		WaitingFor:   wait.ForListeningPort(port),
	}

	if cfg.Network != nil {
		req.Networks = []string{cfg.Network.Name}
		req.NetworkAliases = map[string][]string{
			cfg.Network.Name: {alias},
		}
	}

	container, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
		ContainerRequest: req,
		Started:          true,
	})
	if err != nil {
		t.Fatalf("Failed to start container: %v", err)
	}

	host, err := container.Host(ctx)
	if err != nil {
		t.Fatalf("Failed to extract host: %v", err)
	}
	mappedPort, err := container.MappedPort(ctx, port)
	if err != nil {
		t.Fatalf("Failed to extract port: %v", err)
	}

	t.Setenv("DB_USER", user)
	t.Setenv("DB_PASS", pass)
	t.Setenv("DB_HOST", host)
	t.Setenv("DB_PORT", fmt.Sprint(mappedPort.Int()))
	t.Setenv("DB_NAME", name)
	t.Setenv("DB_NOSSL", "1")

	if cfg.Migrations != nil {
		database.RunMigrations(ctx, cfg.Migrations)
	}

	var pool *pgxpool.Pool
	if !cfg.NoPool {
		pool = database.GetDBPool(ctx)

		err = pool.Ping(ctx)
		if err != nil {
			t.Fatal("Unable to ping database")
		}
	}

	external := &ExternalDatabase{
		Name:     name,
		Password: pass,
		User:     user,
		Port:     port.Int(),
		Host:     alias,
	}

	return &Database{
			Pool:      pool,
			Container: container,
			External:  external,
		}, func() {
			err := container.Terminate(ctx)
			if err != nil {
				t.Error(err)
			}
		}
}