query-orchestration/internal/test/database.go

package test

import (
	"fmt"
	"strconv"
	"sync"
	"testing"

	db "queryorchestration/internal/database"
	"queryorchestration/internal/serviceconfig"

	"github.com/docker/go-connections/nat"
	"github.com/stretchr/testify/require"
	"github.com/testcontainers/testcontainers-go"
	"github.com/testcontainers/testcontainers-go/wait"
)

type CreateDatabaseConfig struct {
	NoMigrations bool
}

const (
	dbAlias = "postgres"
	dbPort  = 5432
)

var dbLock sync.Mutex

func CreateDB(t testing.TB, cfg serviceconfig.ConfigProvider, network string, dcfg *CreateDatabaseConfig) testcontainers.Container {
	port, err := nat.NewPort("tcp", strconv.Itoa(dbPort))
	require.NoError(t, err)

	cfg.SetDBUser("pass")
	cfg.SetDBSecret("postgres")
	cfg.SetDBName(GetAlias(t, "postgres"))
	cfg.SetDBNoSSL(true)

	req := testcontainers.ContainerRequest{
		Image: "postgres:17.2-alpine3.21",
		Name:  "postgres_test_queryorchestration",
		Env: map[string]string{
			"POSTGRES_USER":     cfg.GetDBUser(),
			"POSTGRES_PASSWORD": cfg.GetDBSecret(),
		},
		Cmd:          []string{"postgres", "-c", "max_connections=1000"},
		ExposedPorts: []string{port.Port()},
		WaitingFor: wait.ForAll(
			wait.ForExposedPort(),
			wait.ForListeningPort(port),
			wait.ForLog("database system is ready to accept connections"),
			wait.ForSQL(port, "postgres", func(host string, port nat.Port) string {
				return fmt.Sprintf("postgres://%s:%s@%s:%d/postgres?sslmode=disable",
					cfg.GetDBUser(), cfg.GetDBSecret(), host, port.Int())
			}),
		),
		Networks: []string{network},
		NetworkAliases: map[string][]string{
			network: {dbAlias},
		},
	}

	dbLock.Lock()

	container, err := testcontainers.GenericContainer(t.Context(), testcontainers.GenericContainerRequest{
		ContainerRequest: req,
		Started:          true,
		Reuse:            true,
	})
	require.NoError(t, err)

	host, err := container.Host(t.Context())
	require.NoError(t, err)
	mappedPort, err := container.MappedPort(t.Context(), port)
	require.NoError(t, err)

	dbLock.Unlock()

	cfg.SetDBHost(host)
	cfg.SetDBPort(mappedPort.Int())

	if !dcfg.NoMigrations {
		err := db.RunMigrations(t.Context(), cfg)
		require.NoError(t, err)

		err = cfg.SetDBPool(t.Context())
		require.NoError(t, err)
	}

	return container
}