package test import ( "bufio" "bytes" "encoding/json" "errors" "fmt" "io" "log/slog" "net/http" "strconv" "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 ) func CreateMockServer(t testing.TB) *MockServer { port, err := nat.NewPort("tcp", strconv.Itoa(mockServerPort)) require.NoError(t, err) network := GetNetwork(t) 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}, }, } 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) return &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, } } 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() pageId := 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{}{ { "Id": pageId, "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{} }