query-orchestration/internal/test/database.go

package test

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

	"github.com/docker/go-connections/nat"
	"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.T, ctx context.Context, cfg *CreateDatabaseConfig) (testcontainers.Container, 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.ForAll(
			wait.ForExposedPort(),
			wait.ForListeningPort(port),
			wait.ForLog("database system is ready to accept connections"),
		),
	}

	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)
	}

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

	t.Setenv("DB_USER", dbcfg.DBUser)
	t.Setenv("DB_PASS", dbcfg.DBSecret)
	t.Setenv("DB_NAME", dbcfg.DBName)
	t.Setenv("DB_NOSSL", fmt.Sprintf("%v", dbcfg.DBNoSSL))
	t.Setenv("DB_HOST", dbcfg.DBHost)
	t.Setenv("DB_PORT", fmt.Sprint(dbcfg.DBPort))

	if cfg.Cfg != nil {
		if cfg.RunMigrations {
			cfg.Cfg.SetDBConfig(dbcfg)

			err := migrations.Run(ctx, cfg.Cfg)
			if err != nil {
				t.Fatal(err)
			}

			err = dbcfg.SetDBPool(ctx)
			if err != nil {
				t.Fatal(err)
			}
		}

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

		cfg.Cfg.SetDBConfig(dbcfg)
	}

	return container, func() {
		err := container.Terminate(ctx)
		if err != nil {
			t.Error(err)
		}
	}
}