query-orchestration/internal/test/mockserver.go

package test

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

	"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 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"`
}

const (
	mockServerAlias = "mockserver"
	mockServerPort  = 1080
)

var mockServerLock sync.Mutex

func CreateMockServer(t testing.TB, network string) MockServer {
	port, err := nat.NewPort("tcp", strconv.Itoa(mockServerPort))
	require.NoError(t, err)

	req := testcontainers.ContainerRequest{
		Image:        "mockserver/mockserver:latest",
		Name:         "mockserver_test_queryorchestration",
		ExposedPorts: []string{port.Port()},
		Env: map[string]string{
			"MOCKSERVER_LOG_LEVEL":              "INFO",
			"MOCKSERVER_LIVENESS_HTTP_GET_PATH": "/liveness/probe",
		},
		WaitingFor: wait.ForAll(
			wait.ForExposedPort(),
			wait.ForListeningPort(port),
			wait.ForHTTP("/liveness/probe").
				WithPort(port).
				WithResponseMatcher(func(body io.Reader) bool {
					return true
				}).
				WithStatusCodeMatcher(func(statusCode int) bool {
					return statusCode == http.StatusOK
				}),
		),
		Networks: []string{network},
		NetworkAliases: map[string][]string{
			network: {mockServerAlias},
		},
	}

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

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

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

	mockServerLock.Unlock()

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

	return server
}

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 := 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{}
}