query-orchestration/internal/test/database.go

package test

import (
	"context"
	"fmt"
	"testing"
	"time"

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

	"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 {
	Network       *testcontainers.DockerNetwork
	Cfg           serviceconfig.ConfigProvider
	RunMigrations bool
}

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

	port, err := nat.NewPort("tcp", "5432")
	require.NoError(t, 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.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@localhost:%s/%s?sslmode=disable",
					user, pass, port.Port(), name)
			}).
				WithStartupTimeout(25*time.Second).
				WithPollInterval(10*time.Second),
		),
	}

	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,
	})
	require.NoError(t, err)

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

	dbcfg := &database.DBConfig{
		DBName:   name,
		DBSecret: pass,
		DBUser:   user,
		DBPort:   mappedPort.Int(),
		DBHost:   host,
		DBNoSSL:  true,
	}

	t.Setenv("PGUSER", dbcfg.DBUser)
	t.Setenv("PGPASSWORD", dbcfg.DBSecret)
	t.Setenv("PGDATABASE", dbcfg.DBName)
	t.Setenv("DB_NOSSL", fmt.Sprintf("%v", dbcfg.DBNoSSL))
	t.Setenv("PGHOST", dbcfg.DBHost)
	t.Setenv("PGPORT", fmt.Sprint(dbcfg.DBPort))

	if cfg.Cfg == nil {
		SetCfgProvider(t, cfg.Cfg)
	}

	cfg.Cfg.SetDBConfig(dbcfg)

	if cfg.RunMigrations {
		err := db.RunMigrations(ctx, cfg.Cfg)
		require.NoError(t, err)

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

	if cfg.Network != nil {
		dbcfg.DBHost = alias
		dbcfg.DBPort = port.Int()
		cfg.Cfg.SetDBConfig(dbcfg)
	}

	return container, func() {
		err := container.Terminate(ctx)
		require.NoError(t, err)
	}
}