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
}

type FullDependenciesParams struct {
	NoObjectStore bool
	Runners       *[]RunnerName
}

func CreateFullDependenciesWithParams(t testing.TB, cfg FullDependenciesConfig, params *FullDependenciesParams) 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)

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

		var lock sync.Mutex

		runs := params.Runners
		if runs == nil {
			allRuns := []RunnerName{}
			for _, run := range runners {
				allRuns = append(allRuns, run.Name)
			}
			runs = &allRuns
		}

		for _, runner := range *runs {
			wg.Add(1)
			go func() {
				queue := CreateQueue(t, cfg, runner)

				lock.Lock()
				deps.QueueURLs[runner] = queue
				lock.Unlock()

				wg.Done()
			}()
		}

		if !params.NoObjectStore {
			SetStoreClient(t, t.Context(), cfg, deps.AWSConfig.ExternalEndpoint)
			CreateBucket(t, cfg)
			SetBucketNotifs(t, t.Context(), cfg)
		}

		wg.Done()
	}()

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

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

	wg.Wait()

	return deps
}

func CreateFullDependencies(t testing.TB, ctx context.Context, cfg FullDependenciesConfig) Dependencies {
	return CreateFullDependenciesWithParams(t, cfg, &FullDependenciesParams{})
}

type APINetwork struct {
	Dependencies Dependencies
	API          *Container
}

func CreateAPINetworkWithParams(t testing.TB, cfg FullDependenciesConfig, api API, params *FullDependenciesParams) (*APINetwork, func()) {
	deps := CreateFullDependenciesWithParams(t, cfg, params)

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

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

func CreateAPINetwork(t testing.TB, ctx context.Context, cfg FullDependenciesConfig, api API) (*APINetwork, func()) {
	return CreateAPINetworkWithParams(t, cfg, api, &FullDependenciesParams{})
}

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
}