query-orchestration/internal/test/ecosystem.go

package test

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"strings"
	"testing"
	"time"

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

	"github.com/docker/go-connections/nat"
	"github.com/google/uuid"
	"github.com/stretchr/testify/require"
	"github.com/testcontainers/testcontainers-go"
	"github.com/testcontainers/testcontainers-go/wait"
)

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, clean := CreateFullDependencies(t, ctx, cfg)

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

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

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

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

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, func()) {
	network := GetNetwork(t, ctx)

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

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

	deps.DBConfig = CreateDB(t, ctx, cfg, network, &CreateDatabaseConfig{})
	deps.AWSConfig = CreateAWSContainer(t, ctx, cfg, network)

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

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

	SetBucketNotifs(t, ctx, cfg)

	return deps, func() {
		cleanMock()
	}
}

type APINetwork struct {
	Dependencies Dependencies
	API          *Container
}

func CreateAPINetwork(t testing.TB, ctx context.Context, cfg FullDependenciesConfig, api API) (*APINetwork, func()) {
	deps, depsclean := 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,
		}, func() {
			depsclean()
			ccleanup()
		}
}

type Address string

type MockServer struct {
	Internal  Address
	External  Address
	Container testcontainers.Container
	Client    *http.Client
}

type MockBody any
type MockQueries map[string][]string
type MockHeaders map[string][]string

type MockRequest struct {
	Method  string      `json:"method"`
	Path    string      `json:"path"`
	Headers MockHeaders `json:"headers"`
	Body    MockBody    `json:"body"`
	Query   MockQueries `json:"queryStringParameters"`
}

type MockResponse struct {
	Code    int         `json:"statusCode"`
	Headers MockHeaders `json:"headers"`
	Body    MockBody    `json:"body"`
}

type MockExpectation struct {
	Request  MockRequest  `json:"httpRequest"`
	Response MockResponse `json:"httpResponse"`
}

func CreateMockServer(t testing.TB, ctx context.Context, network string) (MockServer, func()) {
	name := NormaliseAlias(fmt.Sprintf("mockserver_%s", t.Name()))
	port, err := nat.NewPort("tcp", "1080")
	require.NoError(t, err)

	req := testcontainers.ContainerRequest{
		Image:        "mockserver/mockserver:latest",
		ExposedPorts: []string{port.Port()},
		Env: map[string]string{
			"MOCKSERVER_LOG_LEVEL": "INFO",
		},
		WaitingFor: wait.ForAll(
			wait.ForExposedPort(),
			wait.ForListeningPort(port),
		),
		Networks: []string{network},
		NetworkAliases: map[string][]string{
			network: {name},
		},
	}

	container, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
		ContainerRequest: req,
		Started:          true,
	})
	require.NoError(t, err)

	time.Sleep(2 * time.Second)

	host, err := container.Host(ctx)
	require.NoError(t, err)

	externalPort, err := container.MappedPort(ctx, port)
	require.NoError(t, err)

	server := MockServer{
		Client:    &http.Client{},
		Internal:  Address(fmt.Sprintf("http://%s:%d", name, port.Int())),
		External:  Address(fmt.Sprintf("http://%s:%d", host, externalPort.Int())),
		Container: container,
	}

	return server, func() {
		err := container.Terminate(ctx)
		require.NoError(t, err)
	}
}

func NormaliseAlias(fullName string) string {
	name := strings.ToLower(fullName)
	name = strings.ReplaceAll(name, "/", "")
	name = strings.ReplaceAll(name, " ", "")
	name = strings.ReplaceAll(name, "-", "")
	return name
}

func CreateMockExpectation(t testing.TB, server MockServer, expectation MockExpectation) {
	jsonData, err := json.Marshal(expectation)
	require.NoError(t, err)

	url := fmt.Sprintf("%s/mockserver/expectation", server.External)
	req, err := http.NewRequest("PUT", url, bytes.NewBuffer(jsonData))
	require.NoError(t, err)
	req.Header.Set("Content-Type", "application/json")

	resp, err := server.Client.Do(req)
	require.NoError(t, err)
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusCreated && resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		t.Fatalf("failed to configure MockServer, status: %d, response: %s", resp.StatusCode, body)
	}
}

func CreateDetectDocumentTextExpectation(t testing.TB, mockServer MockServer, body string) MockExpectation {
	childId := uuid.NewString()
	expectation := MockExpectation{
		Request: MockRequest{
			Method: "POST",
			Path:   "/",
			Headers: MockHeaders{
				"X-Amz-Target": []string{"Textract.AnalyzeDocument"},
			},
			Body: map[string]interface{}{
				"Document":     map[string]interface{}{},
				"FeatureTypes": []string{"LAYOUT", "SIGNATURES"},
			},
			Query: MockQueries{},
		},
		Response: MockResponse{
			Code: 200,
			Headers: MockHeaders{
				"Content-Type": {"application/json"},
			},
			Body: map[string]interface{}{
				"Blocks": []map[string]interface{}{
					{
						"BlockType": "PAGE",
						"Relationships": []map[string]interface{}{
							{
								"Type": "CHILD",
								"Ids": []string{
									childId,
								},
							},
						},
					},
					{
						"Id":   childId,
						"Text": body,
					},
				},
			},
		},
	}

	CreateMockExpectation(t, mockServer, expectation)

	return expectation
}

func WaitForMockEndpoint(t testing.TB, server MockServer, request MockRequest) MockRequest {
	t.Helper()

	verificationRequest := map[string]any{
		"httpRequest": request,
		"times": map[string]any{
			"atLeast": 1,
		},
	}

	jsonData, err := json.Marshal(verificationRequest)
	require.NoError(t, err)

	verifyURL := fmt.Sprintf("%s/mockserver/verify", server.External)

	clientTimeout := 500 * time.Millisecond
	client := &http.Client{
		Timeout: clientTimeout,
	}

	timeout := time.After(60 * time.Second)
	ticker := time.NewTicker(clientTimeout)
	defer ticker.Stop()

	slog.Info("Attempting to process request", "body", jsonData)

	for {
		select {
		case <-timeout:
			require.NoError(t, errors.New("Timeout waiting for mock http request to be fulfilled"))
		case <-ticker.C:
			req, err := http.NewRequest("PUT", verifyURL, bytes.NewBuffer(jsonData))
			require.NoError(t, err)
			req.Header.Set("Content-Type", "application/json")

			resp, err := client.Do(req)
			require.NoError(t, err)

			scanner := bufio.NewScanner(resp.Body)
			for scanner.Scan() {
				line := scanner.Text()
				slog.Info(line)
			}

			if resp.StatusCode == http.StatusAccepted {
				slog.Info("request found")
				resp.Body.Close()
				return retrieveMatchingRequest(t, server, request)
			}

			resp.Body.Close()
			slog.Error("no request found")
		}
	}
}

func retrieveMatchingRequest(t testing.TB, server MockServer, request MockRequest) MockRequest {
	retrieveURL := fmt.Sprintf("%s/mockserver/retrieve?type=REQUESTS", server.External)

	req, err := http.NewRequest("PUT", retrieveURL, nil)
	require.NoError(t, err)

	client := &http.Client{}
	resp, err := client.Do(req)
	require.NoError(t, err)
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	require.NoError(t, err)

	var requests []MockRequest
	err = json.Unmarshal(body, &requests)
	require.NoError(t, err)

	for i := len(requests) - 1; i >= 0; i-- {
		attemptRequest := requests[i]

		if attemptRequest.Method == request.Method && attemptRequest.Path == request.Path {
			return attemptRequest
		}
	}

	require.Fail(t, "no request found")
	return MockRequest{}
}