query-orchestration/internal/test/ecosystem.go

package test

import (
	"context"
	"os"
	"path"
	"testing"

	"queryorchestration/internal/serviceconfig"
	"queryorchestration/internal/serviceconfig/aws"
	"queryorchestration/internal/serviceconfig/database"
	"queryorchestration/internal/serviceconfig/objectstore"
	queryapi "queryorchestration/pkg/queryAPI"

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

func SetCfgProviderWithBasePath(t testing.TB, cfg serviceconfig.ConfigProvider, basePath string) {
	SetCfgProvider(t, cfg)

	cfg.SetBasePath(path.Join(os.Getenv("PWD"), basePath))
}

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

func CreateAPINetwork(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 := CreateAPI(t, ctx, &APIConfig{
		API:     scfg.API,
		Cfg:     scfg.Cfg,
		Network: network,
	})

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

type FullDependenciesConfig interface {
	serviceconfig.ConfigProvider
	objectstore.ConfigProvider
}

type Dependencies struct {
	BucketName string
	QueueURLs  map[string]string
	Network    *testcontainers.DockerNetwork
	AWSConfig  *AWSContainerConfig
}

func CreateFullDependencies(t testing.TB, ctx context.Context, cfg FullDependenciesConfig) (Dependencies, func()) {
	network, ncleanup := CreateNetwork(t, ctx)

	acfg, awsclean := CreateAWSContainer(t, ctx, &CreateAWSConfig{
		Cfg:     cfg,
		Network: network,
	})

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

	SetQueueClient(t, ctx, cfg)
	SetStoreClient(t, ctx, cfg, acfg.ExternalEndpoint)

	urls := map[string]string{}
	for _, runner := range runners {
		urls[runner.Name] = CreateQueue(t, ctx, cfg, runner.Name)
	}

	CreateBucket(t, ctx, cfg, BucketName)
	SetBucketNotifs(t, ctx, cfg, BucketName)

	return Dependencies{
			BucketName: BucketName,
			QueueURLs:  urls,
			Network:    network,
			AWSConfig:  acfg,
		}, func() {
			awsclean()
			dbcleanup()
			ncleanup()
		}
}

type Network struct {
	Dependencies Dependencies
	APIs         map[string]*Container
	Runners      map[string]*Container
	Client       *queryapi.ClientWithResponses
}

func CreateFullNetwork(t testing.TB, ctx context.Context, cfg FullDependenciesConfig) (Network, func()) {
	deps, clean := CreateFullDependencies(t, ctx, cfg)

	apiContainers := make(map[string]*Container, len(apis))
	apiClean := make([]func(), len(apis))
	for i, s := range apis {
		c, ccleanup := CreateAPI(t, ctx, &APIConfig{
			API:     s,
			Cfg:     cfg,
			Network: deps.Network,
		})
		apiContainers[s.Name] = c
		apiClean[i] = ccleanup
	}

	qService, err := queryapi.NewClientWithResponses(apiContainers[QueryAPIName].URI)
	require.NoError(t, err)

	runnerContainers := make(map[string]*Container, len(runners))
	runnerClean := make([]func(), len(runners))
	for i, r := range runners {
		c, ccleanup := CreateRunner(t, ctx, &RunnerConfig{
			Runner:  r,
			Cfg:     cfg,
			Network: deps.Network,
		})

		runnerContainers[r.Name] = c
		runnerClean[i] = ccleanup
	}

	return Network{
			Dependencies: deps,
			APIs:         apiContainers,
			Runners:      runnerContainers,
			Client:       qService,
		}, func() {
			for _, c := range apiClean {
				c()
			}
			for _, c := range runnerClean {
				c()
			}
			clean()
		}
}