package clients import ( "context" "encoding/json" "io" "net/http" "net/http/httptest" "strings" "testing" "github.com/easyai/easyai-ai-gateway/apps/api/internal/store" ) func TestDecodeHTTPResponseForProtocolCapturesOfficialErrorWire(t *testing.T) { response := &http.Response{ StatusCode: http.StatusTooManyRequests, Status: "429 Too Many Requests", Header: http.Header{ "Content-Type": {"application/json"}, "X-Request-Id": {"req_1"}, "X-Ratelimit-Reset-Requests": {"1s"}, "Set-Cookie": {"secret=must-not-leak"}, "Connection": {"keep-alive"}, }, Body: io.NopCloser(strings.NewReader(`{"error":{"message":"slow down","future_field":true}}`)), } _, wire, err := decodeHTTPResponseForProtocol(response, ProtocolOpenAIResponses) if err == nil { t.Fatal("expected upstream error") } if wire == nil || wire.StatusCode != http.StatusTooManyRequests || wire.Protocol != ProtocolOpenAIResponses { t.Fatalf("unexpected wire response: %+v", wire) } if wire.Headers["X-Request-Id"][0] != "req_1" || wire.Headers["X-Ratelimit-Reset-Requests"][0] != "1s" { t.Fatalf("official response headers were lost: %+v", wire.Headers) } if _, ok := wire.Headers["Set-Cookie"]; ok { t.Fatalf("sensitive header leaked: %+v", wire.Headers) } if _, ok := wire.Headers["Connection"]; ok { t.Fatalf("hop-by-hop header leaked: %+v", wire.Headers) } if ErrorWireResponse(err) != wire { t.Fatal("client error did not retain its wire response") } } func TestGeminiNativeStreamPreservesOfficialEvents(t *testing.T) { var requestedPath string upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { requestedPath = r.URL.RequestURI() w.Header().Set("Content-Type", "text/event-stream") w.Header().Set("X-Goog-Request-Id", "goog-stream-1") _, _ = io.WriteString(w, "data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"hello\"}]}}],\"futureOfficialField\":true}\n\n") })) defer upstream.Close() var events []StreamDeltaEvent response, err := (GeminiClient{HTTPClient: upstream.Client()}).Run(context.Background(), Request{ Kind: "chat.completions", ModelType: "text", Model: "gemini-test", Stream: true, Body: map[string]any{"prompt": "hello"}, Candidate: store.RuntimeModelCandidate{ Provider: "gemini", BaseURL: upstream.URL, ProviderModelName: "gemini-test", Credentials: map[string]any{"apiKey": "test-key"}, }, StreamDelta: func(event StreamDeltaEvent) error { events = append(events, event) return nil }, }) if err != nil { t.Fatalf("Gemini stream failed: %v", err) } if !strings.Contains(requestedPath, ":streamGenerateContent") || !strings.Contains(requestedPath, "alt=sse") { t.Fatalf("unexpected Gemini stream endpoint: %s", requestedPath) } if len(events) != 1 || events[0].WireProtocol != ProtocolGeminiGenerateContent || events[0].Event["futureOfficialField"] != true { raw, _ := json.Marshal(events) t.Fatalf("official Gemini event was not preserved: %s", raw) } if response.Wire == nil || response.Wire.Headers["X-Goog-Request-Id"][0] != "goog-stream-1" { t.Fatalf("Gemini stream wire metadata was lost: %+v", response.Wire) } }