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 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 newGateway := func(calls *atomic.Int64) *httptest.Server { return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { calls.Add(1) 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}, apiKey: "key-1", 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()) } } 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{ apiKey: "key-1", runID: "run-1", runToken: "token-1", gatewayTLSName: "ai.example.com", } opts.setHeaders(request, false) if request.Host != "ai.example.com" { t.Fatalf("Host=%q", request.Host) } }