query-orchestration/internal/test/ecosystem.go

package test

import (
	"context"
	"queryorchestration/internal/serviceconfig"
	"testing"

	"github.com/testcontainers/testcontainers-go"
)

type EcosystemConfig struct {
	Services map[Service]*Container
	Runners  map[Runner]*Container
	Cfg      *serviceconfig.BaseConfig
}

type EcosystemNetworkConfig struct {
	Runners  []*RunnerNetworkConfig
	Services []*ServiceNetworkConfig
	Cfg      *serviceconfig.BaseConfig
	Network  *testcontainers.DockerNetwork
}

func CreateRunnersAndServicesNetwork(t *testing.T, ctx context.Context, ncfg *EcosystemNetworkConfig) (*EcosystemConfig, func()) {
	if ncfg.Cfg == nil {
		ncfg.Cfg = CreateBaseConfig()
	}

	network := ncfg.Network
	var ncleanup func()
	if network == nil {
		network, ncleanup = CreateNetwork(t, ctx)
	}

	_, dbcleanup := CreateDB(t, ctx, &CreateDatabaseConfig{
		Network:       network,
		Cfg:           ncfg.Cfg,
		RunMigrations: true,
	})

	ss := make(map[string]*Container, len(ncfg.Services))
	sclean := make([]func(), len(ncfg.Services))
	for i, s := range ncfg.Services {
		c, ccleanup := CreateService(t, ctx, &ServiceConfig{
			Name:    s.Name,
			Env:     s.Env,
			Cfg:     ncfg.Cfg,
			Network: network,
		})
		ss[s.Name] = c
		sclean[i] = ccleanup
	}

	var qcleanup func()
	if ncfg.Cfg.QueueClient == nil {
		qcleanup = CreateQueueClient(t, ctx, &CreateQueueConfig{
			Network: network,
			Cfg:     ncfg.Cfg,
		})
	}

	qs := make(map[string]*Container, len(ncfg.Runners))
	qclean := make([]func(), len(ncfg.Runners))
	for i, r := range ncfg.Runners {
		url := CreateQueue(t, ctx, ncfg.Cfg, r.Name)

		c, ccleanup := CreateRunner(t, ctx, &RunnerConfig{
			Name:     r.Name,
			Env:      r.Env,
			QueueURL: url,
			Cfg:      ncfg.Cfg,
			Network:  network,
		})

		qs[r.Name] = c
		qclean[i] = ccleanup
	}

	return &EcosystemConfig{
			Services: ss,
			Runners:  qs,
		}, func() {
			for _, c := range sclean {
				c()
			}
			for _, c := range qclean {
				c()
			}
			dbcleanup()
			if qcleanup != nil {
				qcleanup()
			}
			if ncleanup != nil {
				ncleanup()
			}
		}
}

func CreateBaseConfig() *serviceconfig.BaseConfig {
	cfg := &serviceconfig.BaseConfig{}

	cfg.AWSKeyID = "test"
	cfg.AWSSecretKey = "test"
	cfg.AWSRegion = "us-east-1"
	cfg.DBUser = "invalid_user"
	cfg.DBSecret = "invalid_pass"
	cfg.DBHost = "invalid_host"
	cfg.DBPort = 5432
	cfg.DBName = "invalid_name"
	cfg.DBNoSSL = true

	return cfg
}

type ServiceNetworkConfig struct {
	Name Runner
	Env  map[string]string
}

func CreateServiceNetwork(t *testing.T, ctx context.Context, scfg *ServiceNetworkConfig) (*Container, func()) {
	cfg := CreateBaseConfig()

	network, ncleanup := CreateNetwork(t, ctx)

	_, dbcleanup := CreateDB(t, ctx, &CreateDatabaseConfig{
		Network: network,
		Cfg:     cfg,
	})

	c, ccleanup := CreateService(t, ctx, &ServiceConfig{
		Name:    scfg.Name,
		Env:     scfg.Env,
		Cfg:     cfg,
		Network: network,
	})

	return c, func() {
		dbcleanup()
		ccleanup()
		ncleanup()
	}
}

type RunnerNetworkConfig struct {
	Network *testcontainers.DockerNetwork
	Cfg     *serviceconfig.BaseConfig
	Name    Runner
	Env     map[string]string
}

func CreateRunnerNetwork(t *testing.T, ctx context.Context, rcfg *RunnerNetworkConfig) (*Container, func()) {
	if rcfg.Cfg == nil {
		rcfg.Cfg = CreateBaseConfig()
	}

	network := rcfg.Network
	var ncleanup func()
	if network == nil {
		network, ncleanup = CreateNetwork(t, ctx)
	}

	var qcleanup func()
	if rcfg.Cfg.QueueClient == nil {
		qcleanup = CreateQueueClient(t, ctx, &CreateQueueConfig{
			Network: network,
			Cfg:     rcfg.Cfg,
		})
	}

	url := CreateQueue(t, ctx, rcfg.Cfg, rcfg.Name)

	_, dbcleanup := CreateDB(t, ctx, &CreateDatabaseConfig{
		Network: network,
		Cfg:     rcfg.Cfg,
	})

	c, ccleanup := CreateRunner(t, ctx, &RunnerConfig{
		Name:     rcfg.Name,
		Env:      rcfg.Env,
		QueueURL: url,
		Cfg:      rcfg.Cfg,
		Network:  network,
	})

	return c, func() {
		ccleanup()
		dbcleanup()
		if qcleanup != nil {
			qcleanup()
		}
		if ncleanup != nil {
			ncleanup()
		}
	}
}