query-orchestration/internal/test/ecosystem.go

package test

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

	"github.com/stretchr/testify/require"
	"github.com/testcontainers/testcontainers-go"
)

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

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

func CreateRunnersAndServicesNetwork(t testing.TB, ctx context.Context, ncfg *EcosystemNetworkConfig) (*EcosystemConfig, func()) {
	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.AWSContainer == nil {
		_, qcleanup = CreateAWSContainer(t, ctx, &CreateAWSConfig{
			Network: network,
			Cfg:     ncfg.Cfg,
		})
	}

	if ncfg.Cfg.GetQueueClient() == nil {
		err := ncfg.Cfg.SetQueueClient(ctx)
		require.NoError(t, err)
	}

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

		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 SetCfgProvider(t testing.TB, cfg serviceconfig.ConfigProvider) {
	cfg.SetAWSConfig(&aws.AWSConfig{
		AWSKeyID:     "test",
		AWSSecretKey: "test",
		AWSRegion:    "us-east-1",
	})
	t.Setenv("AWS_ACCESS_KEY_ID", cfg.GetAWSKeyID())
	t.Setenv("AWS_SECRET_ACCESS_KEY", cfg.GetAWSSecretKey())
	t.Setenv("AWS_SESSION_TOKEN", cfg.GetAWSSessionToken())
	t.Setenv("AWS_REGION", cfg.GetAWSRegion())
	t.Setenv("AWS_DEFAULT_REGION", cfg.GetAWSRegion())

	cfg.SetDBConfig(&database.DBConfig{
		DBUser:   "invalid_user",
		DBSecret: "invalid_pass",
		DBHost:   "invalid_host",
		DBPort:   5432,
		DBName:   "invalid_name",
		DBNoSSL:  true,
	})
}

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

func CreateServiceNetwork(t testing.TB, ctx context.Context, scfg *ServiceNetworkConfig) (*Container, func()) {
	if scfg.Cfg == nil {
		scfg.Cfg = &serviceconfig.BaseConfig{}
		SetCfgProvider(t, scfg.Cfg)
	}

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

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

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

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

type RunnerNetworkConfig struct {
	Network      *testcontainers.DockerNetwork
	AWSContainer *AWSContainerConfig
	QueueURL     *string
	Cfg          serviceconfig.ConfigProvider
	Name         Runner
	Env          map[string]string
}

func CreateRunnerNetwork(t testing.TB, ctx context.Context, rcfg *RunnerNetworkConfig) (*Container, func()) {
	network := rcfg.Network
	var ncleanup func()
	if network == nil {
		network, ncleanup = CreateNetwork(t, ctx)
	}

	var qcleanup func()
	if rcfg.AWSContainer == nil {
		_, qcleanup = CreateAWSContainer(t, ctx, &CreateAWSConfig{
			Network: network,
			Cfg:     rcfg.Cfg,
		})
	}

	if rcfg.Cfg.GetQueueClient() == nil {
		err := rcfg.Cfg.SetQueueClient(ctx)
		require.NoError(t, err)
	}

	url := rcfg.QueueURL
	if url == nil {
		u := CreateQueue(t, ctx, rcfg.Cfg, rcfg.Name)
		url = &u
	}

	_, 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()
		}
	}
}