query-orchestration/internal/test/container.go

package test

import (
	"context"
	"fmt"
	"strconv"
	"testing"

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

type containerConfig struct {
	ServiceName  string
	DB           *ExternalDatabase
	Network      *testcontainers.DockerNetwork
	Env          map[string]string
	WaitForMsg   string
	ExposedPorts []nat.Port
}

func createContainer(t *testing.T, ctx context.Context, cfg *containerConfig) (testcontainers.Container, func()) {
	env := map[string]string{
		"DB_USER":  cfg.DB.User,
		"DB_PASS":  cfg.DB.Password,
		"DB_HOST":  cfg.DB.Host,
		"DB_NAME":  cfg.DB.Name,
		"DB_PORT":  strconv.Itoa(cfg.DB.Port),
		"DB_NOSSL": "1",
	}
	if cfg.Env != nil {
		for k, v := range cfg.Env {
			env[k] = v
		}
	}

	req := testcontainers.ContainerRequest{
		Image:      "queryorchestration:latest",
		Env:        env,
		Networks:   []string{cfg.Network.Name},
		WaitingFor: wait.ForLog(cfg.WaitForMsg),
		Entrypoint: []string{fmt.Sprintf("./%s", cfg.ServiceName)},
	}

	if len(cfg.ExposedPorts) > 0 {
		ports := make([]string, len(cfg.ExposedPorts))
		for index, port := range cfg.ExposedPorts {
			ports[index] = port.Port()
		}
		req.ExposedPorts = ports
		req.WaitingFor = wait.ForAll(
			// wait.ForExposedPort(),
			// wait.ForListeningPort(cfg.ExposedPorts[0]),
			wait.ForLog(cfg.WaitForMsg),
		)
	}

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

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