easyai-ai-gateway/apps/api/cmd/acceptance-load/main_test.go

package main

import (
	"bytes"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"sync/atomic"
	"testing"
)

func TestStreamGeminiImageHashDoesNotNeedWholeResponse(t *testing.T) {
	payload := paddedPNG(4 << 20)
	response := fmt.Sprintf(`{"candidates":[{"content":{"parts":[{"inlineData":{"mimeType":"image/png","data":"%s"}}]}}]}`,
		base64.StdEncoding.EncodeToString(payload))
	size, digest, err := streamGeminiImageHash(bytes.NewBufferString(response))
	if err != nil {
		t.Fatalf("stream Gemini output: %v", err)
	}
	expected := sha256.Sum256(payload)
	if size != int64(len(payload)) || digest != hex.EncodeToString(expected[:]) {
		t.Fatalf("size=%d digest=%s", size, digest)
	}
}

func TestPaddedPNGVariantsAreExactSizeAndUnique(t *testing.T) {
	first := paddedPNGVariant(256<<10, 1)
	second := paddedPNGVariant(256<<10, 2)
	if len(first) != 256<<10 || len(second) != 256<<10 {
		t.Fatalf("variant sizes=%d/%d", len(first), len(second))
	}
	firstHash := sha256.Sum256(first)
	secondHash := sha256.Sum256(second)
	if firstHash == secondHash {
		t.Fatal("distinct variants have the same SHA-256")
	}
}

func TestStreamGeminiRequestBodyPreservesInput(t *testing.T) {
	input := paddedPNGVariant(2<<20, 37)
	var body struct {
		Contents []struct {
			Parts []struct {
				InlineData *struct {
					MIMEType string `json:"mimeType"`
					Data     string `json:"data"`
				} `json:"inlineData"`
			} `json:"parts"`
		} `json:"contents"`
		GenerationConfig struct {
			ResponseModalities []string `json:"responseModalities"`
		} `json:"generationConfig"`
	}
	stream := streamGeminiRequestBody(input)
	defer stream.Close()
	if err := json.NewDecoder(stream).Decode(&body); err != nil {
		t.Fatalf("decode streamed Gemini request: %v", err)
	}
	if len(body.Contents) != 1 || len(body.Contents[0].Parts) != 2 || body.Contents[0].Parts[1].InlineData == nil {
		t.Fatalf("unexpected Gemini body structure: %+v", body)
	}
	inlineData := body.Contents[0].Parts[1].InlineData
	if inlineData.MIMEType != "image/png" {
		t.Fatalf("mime type=%q", inlineData.MIMEType)
	}
	decoded, err := base64.StdEncoding.DecodeString(inlineData.Data)
	if err != nil {
		t.Fatalf("decode streamed input: %v", err)
	}
	if !bytes.Equal(decoded, input) {
		t.Fatal("streamed input differs from source")
	}
	if len(body.GenerationConfig.ResponseModalities) != 1 || body.GenerationConfig.ResponseModalities[0] != "IMAGE" {
		t.Fatalf("response modalities=%v", body.GenerationConfig.ResponseModalities)
	}
}

func TestVideoCombinationsProvide128UniqueInputs(t *testing.T) {
	images := make([]string, 16)
	for index := range images {
		images[index] = fmt.Sprintf("https://fixtures.example/image-%02d", index)
	}
	combinations := videoCombinations(images, 128)
	if len(combinations) != 128 {
		t.Fatalf("combinations=%d", len(combinations))
	}
	seen := map[string]struct{}{}
	for _, combination := range combinations {
		seen[fmt.Sprint(combination)] = struct{}{}
		if len(combination) != 3 && len(combination) != 6 && len(combination) != 9 {
			t.Fatalf("invalid combination size=%d", len(combination))
		}
	}
	if len(seen) != 128 {
		t.Fatalf("unique combinations=%d", len(seen))
	}
}

func TestGeminiLoadIsSplitAcrossTwoGatewayAPIs(t *testing.T) {
	output := paddedPNG(256 << 10)
	encoded := base64.StdEncoding.EncodeToString(output)
	var firstCalls atomic.Int64
	var secondCalls atomic.Int64
	var firstKeyCalls atomic.Int64
	var secondKeyCalls atomic.Int64
	newGateway := func(calls *atomic.Int64) *httptest.Server {
		return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			calls.Add(1)
			switch r.Header.Get("Authorization") {
			case "Bearer key-1":
				firstKeyCalls.Add(1)
			case "Bearer key-2":
				secondKeyCalls.Add(1)
			default:
				t.Errorf("unexpected authorization header")
			}
			if r.Header.Get(runHeader) != "run-1" || r.Header.Get(tokenHeader) != "token-1" {
				t.Errorf("missing acceptance headers")
			}
			var body map[string]any
			if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
				t.Errorf("decode Gemini body: %v", err)
			}
			_ = json.NewEncoder(w).Encode(map[string]any{
				"candidates": []any{map[string]any{"content": map[string]any{
					"parts": []any{map[string]any{"inlineData": map[string]any{
						"mimeType": "image/png", "data": encoded,
					}}},
				}}},
			})
		}))
	}
	first := newGateway(&firstCalls)
	defer first.Close()
	second := newGateway(&secondCalls)
	defer second.Close()

	opts := options{
		gateways: []string{first.URL, second.URL}, apiKeys: []string{"key-1", "key-2"}, runID: "run-1",
		runToken: "token-1", geminiModel: "gemini-image-test",
	}
	result := runGemini(t.Context(), http.DefaultClient, opts, "dual-api", 8, 256<<10, 256<<10, false)
	if result.err != nil || result.report.Completed != 8 {
		t.Fatalf("result=%+v err=%v", result.report, result.err)
	}
	if firstCalls.Load() != 4 || secondCalls.Load() != 4 {
		t.Fatalf("gateway calls=%d/%d", firstCalls.Load(), secondCalls.Load())
	}
	if firstKeyCalls.Load() != 4 || secondKeyCalls.Load() != 4 {
		t.Fatalf("API key calls=%d/%d", firstKeyCalls.Load(), secondKeyCalls.Load())
	}
}

func TestValidateVideoAssetDownloadsFinalMedia(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set("Content-Type", "video/mp4")
		_, _ = w.Write([]byte{0, 0, 0, 16, 'f', 't', 'y', 'p', 'i', 's', 'o', 'm'})
	}))
	defer server.Close()
	if err := validateVideoAsset(t.Context(), server.Client(), server.URL+"/result.mp4"); err != nil {
		t.Fatalf("validate video asset: %v", err)
	}
	if got := findMediaURL(map[string]any{"content": map[string]any{"video_url": server.URL}}); got != server.URL {
		t.Fatalf("media URL=%q", got)
	}
}

func TestAcceptanceReportErrorRedactsSecretsAndURLs(t *testing.T) {
	got := redactError(
		`token-1 failed at https://example.invalid/video.mp4?token=signed`,
		options{runToken: "token-1"},
	)
	if got != `[REDACTED] failed at [REDACTED_URL]` {
		t.Fatalf("redacted error=%q", got)
	}
}

func TestAcceptanceGatewayTLSNameSetsHostHeader(t *testing.T) {
	request := httptest.NewRequest(http.MethodGet, "https://127.0.0.1/api/v1/healthz", nil)
	opts := options{
		apiKeys: []string{"key-1"}, runID: "run-1", runToken: "token-1",
		gatewayTLSName: "ai.example.com",
	}
	opts.setHeaders(request, 0, false)
	if request.Host != "ai.example.com" {
		t.Fatalf("Host=%q", request.Host)
	}
}