query-orchestration/internal/test/ecosystem.go

package test

import (
	"context"
	"fmt"
	"strings"
	"sync"
	"testing"

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

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

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

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

	var wg sync.WaitGroup

	apiContainers := make(map[APIName]*Container, len(apis))
	apiClean := make([]func(), len(apis))
	var apiMu sync.Mutex
	for i, s := range apis {
		wg.Add(1)
		go func() {
			c, ccleanup := CreateAPI(t, ctx, cfg, deps.Network, &APIConfig{
				API:      s,
				MockHTTP: string(deps.MockServer.Internal),
			})

			apiMu.Lock()
			apiContainers[s.Name] = c
			apiMu.Unlock()

			apiClean[i] = ccleanup
			wg.Done()
		}()
	}

	runnerContainers := make(map[RunnerName]*Container, len(runners))
	runnerClean := make([]func(), len(runners))
	var runnerMu sync.Mutex
	for i, r := range runners {
		wg.Add(1)
		go func() {
			c, ccleanup := CreateRunner(t, ctx, cfg, deps.Network, &RunnerConfig{
				Runner:   r,
				MockHTTP: string(deps.MockServer.Internal),
			})

			runnerMu.Lock()
			runnerContainers[r.Name] = c
			runnerMu.Unlock()

			runnerClean[i] = ccleanup
			wg.Done()
		}()
	}

	wg.Wait()

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

	return Network{
			Dependencies: deps,
			APIs:         apiContainers,
			Runners:      runnerContainers,
			Client:       qService,
		}, func() {
			var wg sync.WaitGroup

			for _, c := range apiClean {
				wg.Add(1)
				go func() {
					c()
					wg.Done()
				}()
			}

			for _, c := range runnerClean {
				wg.Add(1)
				go func() {
					c()
					wg.Done()
				}()
			}

			wg.Wait()
		}
}

type FullDependenciesConfig interface {
	serviceconfig.ConfigProvider
	objectstore.ConfigProvider
}

type Dependencies struct {
	QueueURLs  map[RunnerName]string
	AWSConfig  *AWSContainerConfig
	DBConfig   testcontainers.Container
	MockServer MockServer
	Network    string
}

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

	deps := Dependencies{
		Network:   network,
		QueueURLs: map[RunnerName]string{},
	}

	var wg sync.WaitGroup

	wg.Add(1)
	go func() {
		deps.AWSConfig = CreateAWSContainer(t, cfg, network)

		SetQueueClient(t, ctx, cfg, deps.AWSConfig.ExternalEndpoint)

		for _, runner := range runners {
			wg.Add(1)
			go func() {
				deps.QueueURLs[runner.Name] = CreateQueue(t, ctx, cfg, runner.Name)
				wg.Done()
			}()
		}

		SetStoreClient(t, ctx, cfg, deps.AWSConfig.ExternalEndpoint)
		CreateBucket(t, ctx, cfg)
		SetBucketNotifs(t, ctx, cfg)

		wg.Done()
	}()

	mockServer := CreateMockServer(t, network)
	deps.MockServer = mockServer

	CreateDB(t, cfg, network, &CreateDatabaseConfig{
		NoMigrations: true,
	})

	wg.Wait()

	return deps
}

type APINetwork struct {
	Dependencies Dependencies
	API          *Container
}

func CreateAPINetwork(t testing.TB, ctx context.Context, cfg FullDependenciesConfig, api API) (*APINetwork, func()) {
	deps := CreateFullDependencies(t, ctx, cfg)

	c, ccleanup := CreateAPI(t, ctx, cfg, deps.Network, &APIConfig{
		API:      api,
		MockHTTP: string(deps.MockServer.Internal),
	})

	return &APINetwork{
		Dependencies: deps,
		API:          c,
	}, ccleanup
}

func GetAlias(t testing.TB, baseName string) string {
	name := fmt.Sprintf("%s_%s", baseName, t.Name())
	name = strings.ToLower(name)
	name = strings.ReplaceAll(name, "/", "")
	name = strings.ReplaceAll(name, " ", "")
	name = strings.ReplaceAll(name, "-", "")
	name = strings.ReplaceAll(name, "_", "")
	return name
}