mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-12 04:54:45 +00:00
Model discovery used each API's default page size: 20 for Anthropic and 50 for Gemini. A full catalog needed several sequential requests, and at pod start that overran the 5s warm-fill budget. One replica then logged "initial discovery ... context deadline exceeded" and went without that provider's models until the next probe. Every page now asks for the documented maximum of 1000: Anthropic limit and Gemini pageSize. Gemini gets pageSize on each page, because it requires the same parameters across one listing. Claude-Session: https://claude.ai/code/session_01V5jYJNxpDFdJkyyExpADrv
774 lines
28 KiB
Go
774 lines
28 KiB
Go
package traefikllmgateway
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
// TestNewGeminiAdapter_KeylessIsAConstructorError proves an empty resolved
|
|
// API key is a constructor error for gemini-type providers, mirroring
|
|
// ruling (c) and anthropic's equivalent test.
|
|
func TestNewGeminiAdapter_KeylessIsAConstructorError(t *testing.T) {
|
|
if _, err := newGeminiAdapter("p1", "https://generativelanguage.googleapis.com", ""); err == nil {
|
|
t.Fatalf("newGeminiAdapter: want error for empty API key, got nil")
|
|
}
|
|
}
|
|
|
|
// TestGeminiAdapter_InjectAuth proves injectAuth sets the x-goog-api-key
|
|
// header unconditionally.
|
|
func TestGeminiAdapter_InjectAuth(t *testing.T) {
|
|
a, err := newGeminiAdapter("p1", "https://generativelanguage.googleapis.com", "AIza-test")
|
|
if err != nil {
|
|
t.Fatalf("newGeminiAdapter: %v", err)
|
|
}
|
|
req := httptest.NewRequest(http.MethodPost, "/", nil)
|
|
a.injectAuth(req)
|
|
|
|
if got := req.Header.Get("x-goog-api-key"); got != "AIza-test" { //nolint:gosec // test-only literal, not a real credential
|
|
t.Errorf("x-goog-api-key = %q, want %q", got, "AIza-test")
|
|
}
|
|
}
|
|
|
|
// TestGeminiAdapter_ChatCompletion_NonStreaming proves the adapter builds
|
|
// the {base}/v1beta/models/{model}:generateContent URL from req["model"],
|
|
// sends the x-goog-api-key header, translates the OpenAI request to
|
|
// Gemini's wire format, and translates the Gemini response back to an
|
|
// OpenAI-shaped chat.completion body on the wire, with usage extracted
|
|
// from usageMetadata.
|
|
func TestGeminiAdapter_ChatCompletion_NonStreaming(t *testing.T) {
|
|
const geminiResp = `{"candidates":[{"content":{"role":"model","parts":[{"text":"hi there"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":4},"responseId":"resp_01ABC"}`
|
|
|
|
var gotBody []byte
|
|
var gotAPIKey, gotPath, gotMethod string
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
gotPath = r.URL.Path
|
|
gotMethod = r.Method
|
|
gotBody, _ = io.ReadAll(r.Body)
|
|
gotAPIKey = r.Header.Get("x-goog-api-key")
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(http.StatusOK)
|
|
_, _ = w.Write([]byte(geminiResp))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
a, err := newGeminiAdapter("p1", srv.URL, "AIza-test")
|
|
if err != nil {
|
|
t.Fatalf("newGeminiAdapter: %v", err)
|
|
}
|
|
rec := httptest.NewRecorder()
|
|
|
|
req := map[string]any{
|
|
"model": "gemini-2.5-pro",
|
|
"messages": []any{map[string]any{"role": "user", "content": "hello"}},
|
|
}
|
|
u, err := a.chatCompletion(context.Background(), rec, req)
|
|
if err != nil {
|
|
t.Fatalf("chatCompletion: %v", err)
|
|
}
|
|
|
|
if gotPath != "/v1beta/models/gemini-2.5-pro:generateContent" {
|
|
t.Errorf("path = %q, want /v1beta/models/gemini-2.5-pro:generateContent", gotPath)
|
|
}
|
|
if gotMethod != http.MethodPost {
|
|
t.Errorf("method = %q, want POST", gotMethod)
|
|
}
|
|
if u.prompt != 10 || u.completion != 4 {
|
|
t.Errorf("usage = %+v, want {prompt:10 completion:4}", u)
|
|
}
|
|
if gotAPIKey != "AIza-test" { //nolint:gosec // test-only literal, not a real credential
|
|
t.Errorf("x-goog-api-key = %q, want %q", gotAPIKey, "AIza-test")
|
|
}
|
|
|
|
sent := decodeJSONBody(t, gotBody)
|
|
if _, hasModel := sent["model"]; hasModel {
|
|
t.Errorf("upstream body carries a \"model\" field = %v, want none (Gemini's model is URL-only)", sent["model"])
|
|
}
|
|
contents, ok := sent["contents"].([]any)
|
|
if !ok || len(contents) != 1 {
|
|
t.Fatalf("upstream body contents = %v, want a single-element array", sent["contents"])
|
|
}
|
|
|
|
if got := rec.Header().Get("Content-Type"); got != "application/json" {
|
|
t.Errorf("Content-Type = %q, want application/json", got)
|
|
}
|
|
var out map[string]any
|
|
if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil {
|
|
t.Fatalf("client body is not valid JSON: %v (%s)", err, rec.Body.String())
|
|
}
|
|
if out["object"] != "chat.completion" {
|
|
t.Errorf("object = %v, want %q", out["object"], "chat.completion")
|
|
}
|
|
if out["id"] != "chatcmpl-resp_01ABC" {
|
|
t.Errorf("id = %v, want %q", out["id"], "chatcmpl-resp_01ABC")
|
|
}
|
|
choices, ok := out["choices"].([]any)
|
|
if !ok || len(choices) != 1 {
|
|
t.Fatalf("choices = %v, want a single-element array", out["choices"])
|
|
}
|
|
choice := choices[0].(map[string]any)
|
|
message := choice["message"].(map[string]any)
|
|
if message["content"] != "hi there" {
|
|
t.Errorf("message.content = %v, want %q", message["content"], "hi there")
|
|
}
|
|
if choice["finish_reason"] != "stop" {
|
|
t.Errorf("finish_reason = %v, want %q", choice["finish_reason"], "stop")
|
|
}
|
|
}
|
|
|
|
// TestGeminiAdapter_ChatCompletion_ModelsPrefixStripped proves a
|
|
// req["model"] that already carries a "models/" prefix is not doubled into
|
|
// the URL.
|
|
func TestGeminiAdapter_ChatCompletion_ModelsPrefixStripped(t *testing.T) {
|
|
var gotPath string
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
gotPath = r.URL.Path
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(http.StatusOK)
|
|
_, _ = w.Write([]byte(`{"candidates":[{"content":{"role":"model","parts":[{"text":"ok"}]},"finishReason":"STOP"}]}`))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
a, err := newGeminiAdapter("p1", srv.URL, "AIza-test")
|
|
if err != nil {
|
|
t.Fatalf("newGeminiAdapter: %v", err)
|
|
}
|
|
rec := httptest.NewRecorder()
|
|
|
|
_, err = a.chatCompletion(context.Background(), rec, map[string]any{
|
|
"model": "models/gemini-2.5-pro",
|
|
"messages": []any{map[string]any{"role": "user", "content": "hi"}},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("chatCompletion: %v", err)
|
|
}
|
|
if gotPath != "/v1beta/models/gemini-2.5-pro:generateContent" {
|
|
t.Errorf("path = %q, want /v1beta/models/gemini-2.5-pro:generateContent (no doubled models/ prefix)", gotPath)
|
|
}
|
|
}
|
|
|
|
// TestGeminiAdapter_ChatCompletion_ModelEscaped proves a model id
|
|
// containing "/" is percent-escaped into the upstream URL rather than
|
|
// smuggling an extra path segment in unescaped — review fix (item 2). Go's
|
|
// net/http server auto-decodes r.URL.Path, which would make an escaped and
|
|
// an unescaped "/" indistinguishable, so this asserts on r.URL.EscapedPath()
|
|
// (and cross-checks RequestURI), the wire form the upstream actually
|
|
// receives.
|
|
func TestGeminiAdapter_ChatCompletion_ModelEscaped(t *testing.T) {
|
|
var gotEscapedPath, gotRequestURI string
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
gotEscapedPath = r.URL.EscapedPath()
|
|
gotRequestURI = r.RequestURI
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(http.StatusOK)
|
|
_, _ = w.Write([]byte(`{"candidates":[{"content":{"role":"model","parts":[{"text":"ok"}]},"finishReason":"STOP"}]}`))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
a, err := newGeminiAdapter("p1", srv.URL, "AIza-test")
|
|
if err != nil {
|
|
t.Fatalf("newGeminiAdapter: %v", err)
|
|
}
|
|
rec := httptest.NewRecorder()
|
|
|
|
_, err = a.chatCompletion(context.Background(), rec, map[string]any{
|
|
"model": "org/gemini-2.5-pro",
|
|
"messages": []any{map[string]any{"role": "user", "content": "hi"}},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("chatCompletion: %v", err)
|
|
}
|
|
|
|
const wantEscapedPath = "/v1beta/models/org%2Fgemini-2.5-pro:generateContent"
|
|
if gotEscapedPath != wantEscapedPath {
|
|
t.Errorf("EscapedPath() = %q, want %q (model id's \"/\" must stay percent-escaped, not become an extra path segment)", gotEscapedPath, wantEscapedPath)
|
|
}
|
|
if !strings.Contains(gotRequestURI, "org%2Fgemini-2.5-pro") {
|
|
t.Errorf("RequestURI = %q, want it to contain the escaped model id %q", gotRequestURI, "org%2Fgemini-2.5-pro")
|
|
}
|
|
}
|
|
|
|
// TestGeminiAdapter_Embeddings_ModelEscaped mirrors the chatCompletion
|
|
// escaping test for the embeddings :embedContent URL.
|
|
func TestGeminiAdapter_Embeddings_ModelEscaped(t *testing.T) {
|
|
var gotEscapedPath string
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
gotEscapedPath = r.URL.EscapedPath()
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(http.StatusOK)
|
|
_, _ = w.Write([]byte(`{"embedding":{"values":[0.1]}}`))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
a, err := newGeminiAdapter("p1", srv.URL, "AIza-test")
|
|
if err != nil {
|
|
t.Fatalf("newGeminiAdapter: %v", err)
|
|
}
|
|
rec := httptest.NewRecorder()
|
|
|
|
_, err = a.embeddings(context.Background(), rec, map[string]any{
|
|
"model": "org/embed-004",
|
|
"input": "hi",
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("embeddings: %v", err)
|
|
}
|
|
|
|
const wantEscapedPath = "/v1beta/models/org%2Fembed-004:embedContent"
|
|
if gotEscapedPath != wantEscapedPath {
|
|
t.Errorf("EscapedPath() = %q, want %q", gotEscapedPath, wantEscapedPath)
|
|
}
|
|
}
|
|
|
|
// geminiStreamFrames are the raw SSE "data:" lines a fake streaming
|
|
// upstream emits for the tests below: a role/text chunk, a text-only
|
|
// continuation, then a finish chunk carrying usageMetadata. Gemini's SSE
|
|
// stream has no "event:" field to set, unlike Anthropic's.
|
|
var geminiStreamFrames = []string{
|
|
`{"candidates":[{"content":{"role":"model","parts":[{"text":"Hi!"}]}}],"responseId":"resp_STREAM"}`,
|
|
`{"candidates":[{"content":{"role":"model","parts":[]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":6,"candidatesTokenCount":3}}`,
|
|
}
|
|
|
|
// TestGeminiAdapter_ChatCompletion_Streaming proves the adapter requests
|
|
// the :streamGenerateContent?alt=sse URL, forwards a streaming Gemini
|
|
// response as OpenAI-shaped SSE chunks on the wire, flushing per event, and
|
|
// terminates with "[DONE]" once the upstream stream ends — mirroring
|
|
// ruling (b): there is no distinct terminal event of Gemini's own to key
|
|
// off of.
|
|
func TestGeminiAdapter_ChatCompletion_Streaming(t *testing.T) {
|
|
var gotPath, gotQuery, gotAccept string
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
gotPath = r.URL.Path
|
|
gotQuery = r.URL.RawQuery
|
|
gotAccept = r.Header.Get("Accept")
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
w.WriteHeader(http.StatusOK)
|
|
f, ok := w.(http.Flusher)
|
|
if !ok {
|
|
t.Errorf("httptest ResponseWriter does not implement http.Flusher")
|
|
return
|
|
}
|
|
for _, data := range geminiStreamFrames {
|
|
fmt.Fprintf(w, "data: %s\n\n", data)
|
|
f.Flush()
|
|
}
|
|
}))
|
|
defer srv.Close()
|
|
|
|
a, err := newGeminiAdapter("p1", srv.URL, "AIza-test")
|
|
if err != nil {
|
|
t.Fatalf("newGeminiAdapter: %v", err)
|
|
}
|
|
rec := newRecordingResponseWriter()
|
|
|
|
req := map[string]any{
|
|
"model": "gemini-2.5-pro",
|
|
"stream": true,
|
|
"messages": []any{map[string]any{"role": "user", "content": "hello"}},
|
|
}
|
|
u, err := a.chatCompletion(context.Background(), rec, req)
|
|
if err != nil {
|
|
t.Fatalf("chatCompletion: %v", err)
|
|
}
|
|
|
|
if gotPath != "/v1beta/models/gemini-2.5-pro:streamGenerateContent" {
|
|
t.Errorf("path = %q, want /v1beta/models/gemini-2.5-pro:streamGenerateContent", gotPath)
|
|
}
|
|
if gotQuery != "alt=sse" {
|
|
t.Errorf("query = %q, want %q", gotQuery, "alt=sse")
|
|
}
|
|
if gotAccept != "text/event-stream" {
|
|
t.Errorf("Accept = %q, want text/event-stream", gotAccept)
|
|
}
|
|
if u.prompt != 6 || u.completion != 3 {
|
|
t.Errorf("usage = %+v, want {prompt:6 completion:3}", u)
|
|
}
|
|
|
|
body := rec.body.String()
|
|
if got := rec.header.Get("Content-Type"); got != "text/event-stream" {
|
|
t.Errorf("Content-Type = %q, want text/event-stream", got)
|
|
}
|
|
|
|
wantSubstrings := []string{
|
|
`"delta":{"role":"assistant"}`,
|
|
`"delta":{"content":"Hi!"}`,
|
|
`"finish_reason":"stop"`,
|
|
"data: [DONE]\n\n",
|
|
}
|
|
for _, want := range wantSubstrings {
|
|
if !strings.Contains(body, want) {
|
|
t.Errorf("streamed body missing %q; got:\n%s", want, body)
|
|
}
|
|
}
|
|
// 3 chunks (role, content delta, finish) plus one more flush for the
|
|
// terminal [DONE]: 4 flushes, proving the adapter flushes per event
|
|
// rather than buffering the whole stream.
|
|
const wantFlushes = 4
|
|
if rec.flushed < wantFlushes {
|
|
t.Errorf("flushed = %d, want >= %d", rec.flushed, wantFlushes)
|
|
}
|
|
}
|
|
|
|
// TestGeminiAdapter_ChatCompletion_Streaming_IncludeUsage_EmitsFinalUsageChunk
|
|
// is the M1 regression test: a client that sets
|
|
// stream_options.include_usage on a request routed to a gemini-type
|
|
// provider must still get a final usage-only chunk on the wire, the same
|
|
// way an OpenAI-type provider's own forwardStream relays it from a real
|
|
// upstream usage chunk (provider_openai.go) — Gemini's own stream has no
|
|
// such chunk, so the adapter must synthesize one from the
|
|
// last-usageMetadata-wins total it already accumulated.
|
|
func TestGeminiAdapter_ChatCompletion_Streaming_IncludeUsage_EmitsFinalUsageChunk(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
w.WriteHeader(http.StatusOK)
|
|
f, ok := w.(http.Flusher)
|
|
if !ok {
|
|
t.Fatalf("httptest ResponseWriter does not implement http.Flusher")
|
|
}
|
|
for _, data := range geminiStreamFrames {
|
|
fmt.Fprintf(w, "data: %s\n\n", data)
|
|
f.Flush()
|
|
}
|
|
}))
|
|
defer srv.Close()
|
|
|
|
a, err := newGeminiAdapter("p1", srv.URL, "AIza-test")
|
|
if err != nil {
|
|
t.Fatalf("newGeminiAdapter: %v", err)
|
|
}
|
|
rec := newRecordingResponseWriter()
|
|
|
|
req := map[string]any{
|
|
"model": "gemini-2.5-pro",
|
|
"stream": true,
|
|
"stream_options": map[string]any{"include_usage": true},
|
|
"messages": []any{map[string]any{"role": "user", "content": "hello"}},
|
|
}
|
|
u, err := a.chatCompletion(context.Background(), rec, req)
|
|
if err != nil {
|
|
t.Fatalf("chatCompletion: %v", err)
|
|
}
|
|
if u.prompt != 6 || u.completion != 3 {
|
|
t.Fatalf("usage = %+v, want {prompt:6 completion:3}", u)
|
|
}
|
|
|
|
chunks := sseDataChunks(t, rec.body.String())
|
|
if len(chunks) < 2 {
|
|
t.Fatalf("got %d SSE frames, want at least a usage chunk followed by [DONE]", len(chunks))
|
|
}
|
|
if chunks[len(chunks)-1] != "[DONE]" {
|
|
t.Fatalf("last SSE frame = %q, want [DONE]", chunks[len(chunks)-1])
|
|
}
|
|
usageFrame := chunks[len(chunks)-2]
|
|
var got struct {
|
|
Usage *struct {
|
|
PromptTokens int64 `json:"prompt_tokens"`
|
|
CompletionTokens int64 `json:"completion_tokens"`
|
|
TotalTokens int64 `json:"total_tokens"`
|
|
} `json:"usage"`
|
|
Choices []any `json:"choices"`
|
|
}
|
|
if err := json.Unmarshal([]byte(usageFrame), &got); err != nil {
|
|
t.Fatalf("decode usage chunk %q: %v", usageFrame, err)
|
|
}
|
|
if len(got.Choices) != 0 {
|
|
t.Errorf("usage chunk choices = %v, want empty", got.Choices)
|
|
}
|
|
if got.Usage == nil {
|
|
t.Fatalf("usage chunk has no usage object: %q", usageFrame)
|
|
}
|
|
if got.Usage.PromptTokens != 6 || got.Usage.CompletionTokens != 3 || got.Usage.TotalTokens != 9 {
|
|
t.Errorf("usage chunk = %+v, want {prompt_tokens:6 completion_tokens:3 total_tokens:9}", *got.Usage)
|
|
}
|
|
}
|
|
|
|
// TestGeminiAdapter_ChatCompletion_Streaming_NoIncludeUsage_NoFinalUsageChunk
|
|
// proves the M1 fix is opt-in: a client that never sets
|
|
// stream_options.include_usage gets the stream unchanged from before the
|
|
// fix — no extra usage-only chunk on the wire.
|
|
func TestGeminiAdapter_ChatCompletion_Streaming_NoIncludeUsage_NoFinalUsageChunk(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
w.WriteHeader(http.StatusOK)
|
|
f, ok := w.(http.Flusher)
|
|
if !ok {
|
|
t.Fatalf("httptest ResponseWriter does not implement http.Flusher")
|
|
}
|
|
for _, data := range geminiStreamFrames {
|
|
fmt.Fprintf(w, "data: %s\n\n", data)
|
|
f.Flush()
|
|
}
|
|
}))
|
|
defer srv.Close()
|
|
|
|
a, err := newGeminiAdapter("p1", srv.URL, "AIza-test")
|
|
if err != nil {
|
|
t.Fatalf("newGeminiAdapter: %v", err)
|
|
}
|
|
rec := newRecordingResponseWriter()
|
|
|
|
req := map[string]any{
|
|
"model": "gemini-2.5-pro",
|
|
"stream": true,
|
|
"messages": []any{map[string]any{"role": "user", "content": "hello"}},
|
|
}
|
|
if _, err := a.chatCompletion(context.Background(), rec, req); err != nil {
|
|
t.Fatalf("chatCompletion: %v", err)
|
|
}
|
|
|
|
for _, chunk := range sseDataChunks(t, rec.body.String()) {
|
|
if strings.Contains(chunk, `"usage"`) {
|
|
t.Errorf("SSE frame %q carries a usage field though the client never set stream_options.include_usage", chunk)
|
|
}
|
|
}
|
|
}
|
|
|
|
// TestGeminiAdapter_ChatCompletion_Streaming_CaseInsensitiveContentType
|
|
// proves the streaming-vs-JSON-fallback gate matches an upstream's
|
|
// Content-Type case-insensitively, per lesson (g) — tested with
|
|
// "Text/Event-Stream; charset=utf-8".
|
|
func TestGeminiAdapter_ChatCompletion_Streaming_CaseInsensitiveContentType(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "Text/Event-Stream; charset=utf-8")
|
|
w.WriteHeader(http.StatusOK)
|
|
f, ok := w.(http.Flusher)
|
|
if !ok {
|
|
t.Errorf("httptest ResponseWriter does not implement http.Flusher")
|
|
return
|
|
}
|
|
for _, data := range geminiStreamFrames {
|
|
fmt.Fprintf(w, "data: %s\n\n", data)
|
|
f.Flush()
|
|
}
|
|
}))
|
|
defer srv.Close()
|
|
|
|
a, err := newGeminiAdapter("p1", srv.URL, "AIza-test")
|
|
if err != nil {
|
|
t.Fatalf("newGeminiAdapter: %v", err)
|
|
}
|
|
rec := newRecordingResponseWriter()
|
|
|
|
req := map[string]any{
|
|
"model": "gemini-2.5-pro",
|
|
"stream": true,
|
|
"messages": []any{map[string]any{"role": "user", "content": "hello"}},
|
|
}
|
|
u, err := a.chatCompletion(context.Background(), rec, req)
|
|
if err != nil {
|
|
t.Fatalf("chatCompletion: %v", err)
|
|
}
|
|
if u.prompt != 6 || u.completion != 3 {
|
|
t.Errorf("usage = %+v, want {prompt:6 completion:3} (proves the streaming path, not JSON fallback, was taken)", u)
|
|
}
|
|
if got := rec.header.Get("Content-Type"); got != "text/event-stream" {
|
|
t.Errorf("Content-Type on client response = %q, want %q (sseWriter's own header; JSON fallback would copy the upstream's mixed-case value verbatim)", got, "text/event-stream")
|
|
}
|
|
}
|
|
|
|
// TestGeminiAdapter_ChatCompletion_Upstream429 proves a non-2xx upstream
|
|
// response returns a *providerHTTPError with nothing written to the
|
|
// client.
|
|
func TestGeminiAdapter_ChatCompletion_Upstream429(t *testing.T) {
|
|
const errBody = `{"error":{"code":429,"message":"rate limited","status":"RESOURCE_EXHAUSTED"}}`
|
|
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(http.StatusTooManyRequests)
|
|
_, _ = w.Write([]byte(errBody))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
a, err := newGeminiAdapter("p1", srv.URL, "AIza-test")
|
|
if err != nil {
|
|
t.Fatalf("newGeminiAdapter: %v", err)
|
|
}
|
|
rec := httptest.NewRecorder()
|
|
|
|
u, err := a.chatCompletion(context.Background(), rec, map[string]any{
|
|
"model": "gemini-2.5-pro",
|
|
"messages": []any{map[string]any{"role": "user", "content": "hi"}},
|
|
})
|
|
if err == nil {
|
|
t.Fatalf("chatCompletion: want error, got nil (usage=%+v)", u)
|
|
}
|
|
var httpErr *providerHTTPError
|
|
if !errors.As(err, &httpErr) {
|
|
t.Fatalf("err = %v (%T), want *providerHTTPError", err, err)
|
|
}
|
|
if httpErr.status != http.StatusTooManyRequests {
|
|
t.Errorf("status = %d, want %d", httpErr.status, http.StatusTooManyRequests)
|
|
}
|
|
if rec.Body.Len() != 0 {
|
|
t.Errorf("client response body = %q, want empty", rec.Body.String())
|
|
}
|
|
}
|
|
|
|
// TestGeminiAdapter_ChatCompletion_TranslateErrorSkipsUpstream proves a
|
|
// request field geminiRequestFromOpenAI rejects (n>1, here) never reaches
|
|
// the upstream at all.
|
|
func TestGeminiAdapter_ChatCompletion_TranslateErrorSkipsUpstream(t *testing.T) {
|
|
called := false
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
called = true
|
|
w.WriteHeader(http.StatusOK)
|
|
}))
|
|
defer srv.Close()
|
|
|
|
a, err := newGeminiAdapter("p1", srv.URL, "AIza-test")
|
|
if err != nil {
|
|
t.Fatalf("newGeminiAdapter: %v", err)
|
|
}
|
|
rec := httptest.NewRecorder()
|
|
|
|
_, err = a.chatCompletion(context.Background(), rec, map[string]any{
|
|
"model": "gemini-2.5-pro",
|
|
"messages": []any{map[string]any{"role": "user", "content": "hi"}},
|
|
"n": float64(2),
|
|
})
|
|
var terr *translateError
|
|
if !errors.As(err, &terr) {
|
|
t.Fatalf("err = %v (%T), want *translateError", err, err)
|
|
}
|
|
if called {
|
|
t.Errorf("upstream was called, want the translate error to short-circuit before any request")
|
|
}
|
|
}
|
|
|
|
// TestGeminiAdapter_Embeddings_Single proves embeddings builds the
|
|
// :embedContent URL and translates the response to an OpenAI embeddings
|
|
// list body for a single-string "input".
|
|
func TestGeminiAdapter_Embeddings_Single(t *testing.T) {
|
|
var gotPath string
|
|
var gotBody []byte
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
gotPath = r.URL.Path
|
|
gotBody, _ = io.ReadAll(r.Body)
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(http.StatusOK)
|
|
_, _ = w.Write([]byte(`{"embedding":{"values":[0.1,0.2,0.3]}}`))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
a, err := newGeminiAdapter("p1", srv.URL, "AIza-test")
|
|
if err != nil {
|
|
t.Fatalf("newGeminiAdapter: %v", err)
|
|
}
|
|
rec := httptest.NewRecorder()
|
|
|
|
u, err := a.embeddings(context.Background(), rec, map[string]any{
|
|
"model": "text-embedding-004",
|
|
"input": "hello world",
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("embeddings: %v", err)
|
|
}
|
|
if gotPath != "/v1beta/models/text-embedding-004:embedContent" {
|
|
t.Errorf("path = %q, want /v1beta/models/text-embedding-004:embedContent", gotPath)
|
|
}
|
|
sent := decodeJSONBody(t, gotBody)
|
|
if _, hasRequests := sent["requests"]; hasRequests {
|
|
t.Errorf("single-input body carries \"requests\", want the :embedContent shape")
|
|
}
|
|
if !u.estimated {
|
|
t.Errorf("usage.estimated = false, want true (Gemini reports no embeddings usage)")
|
|
}
|
|
|
|
var out map[string]any
|
|
if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil {
|
|
t.Fatalf("client body is not valid JSON: %v (%s)", err, rec.Body.String())
|
|
}
|
|
if out["object"] != "list" {
|
|
t.Errorf("object = %v, want %q", out["object"], "list")
|
|
}
|
|
data, ok := out["data"].([]any)
|
|
if !ok || len(data) != 1 {
|
|
t.Fatalf("data = %v, want a single-element array", out["data"])
|
|
}
|
|
}
|
|
|
|
// TestGeminiAdapter_Embeddings_Batch proves an array "input" selects
|
|
// Gemini's :batchEmbedContents endpoint and translates its response to an
|
|
// OpenAI embeddings list body with one entry per input string.
|
|
func TestGeminiAdapter_Embeddings_Batch(t *testing.T) {
|
|
var gotPath string
|
|
var gotBody []byte
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
gotPath = r.URL.Path
|
|
gotBody, _ = io.ReadAll(r.Body)
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(http.StatusOK)
|
|
_, _ = w.Write([]byte(`{"embeddings":[{"values":[0.1,0.2]},{"values":[0.3,0.4]}]}`))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
a, err := newGeminiAdapter("p1", srv.URL, "AIza-test")
|
|
if err != nil {
|
|
t.Fatalf("newGeminiAdapter: %v", err)
|
|
}
|
|
rec := httptest.NewRecorder()
|
|
|
|
_, err = a.embeddings(context.Background(), rec, map[string]any{
|
|
"model": "text-embedding-004",
|
|
"input": []any{"hello", "world!!"},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("embeddings: %v", err)
|
|
}
|
|
if gotPath != "/v1beta/models/text-embedding-004:batchEmbedContents" {
|
|
t.Errorf("path = %q, want /v1beta/models/text-embedding-004:batchEmbedContents", gotPath)
|
|
}
|
|
sent := decodeJSONBody(t, gotBody)
|
|
requests, ok := sent["requests"].([]any)
|
|
if !ok || len(requests) != 2 {
|
|
t.Fatalf("requests = %v, want a 2-element array", sent["requests"])
|
|
}
|
|
first := requests[0].(map[string]any)
|
|
if first["model"] != "models/text-embedding-004" {
|
|
t.Errorf("requests[0].model = %v, want %q", first["model"], "models/text-embedding-004")
|
|
}
|
|
|
|
var out map[string]any
|
|
if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil {
|
|
t.Fatalf("client body is not valid JSON: %v (%s)", err, rec.Body.String())
|
|
}
|
|
data, ok := out["data"].([]any)
|
|
if !ok || len(data) != 2 {
|
|
t.Fatalf("data = %v, want a 2-element array", out["data"])
|
|
}
|
|
}
|
|
|
|
// TestGeminiAdapter_ListModels proves listModels parses Gemini's
|
|
// "models": [{"name": "models/..."}] Models API response into a flat
|
|
// slice of bare ids, stripping the "models/" prefix.
|
|
func TestGeminiAdapter_ListModels(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path != "/v1beta/models" {
|
|
t.Errorf("path = %q, want /v1beta/models", r.URL.Path)
|
|
return
|
|
}
|
|
if r.Method != http.MethodGet {
|
|
t.Errorf("method = %q, want GET", r.Method)
|
|
return
|
|
}
|
|
if got := r.Header.Get("x-goog-api-key"); got != "AIza-test" {
|
|
t.Errorf("x-goog-api-key = %q, want %q", got, "AIza-test")
|
|
return
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = w.Write([]byte(`{"models":[{"name":"models/gemini-2.5-pro"},{"name":"models/gemini-2.5-flash"}]}`))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
a, err := newGeminiAdapter("p1", srv.URL, "AIza-test")
|
|
if err != nil {
|
|
t.Fatalf("newGeminiAdapter: %v", err)
|
|
}
|
|
|
|
got, err := a.listModels(context.Background())
|
|
if err != nil {
|
|
t.Fatalf("listModels: %v", err)
|
|
}
|
|
want := []string{"gemini-2.5-pro", "gemini-2.5-flash"}
|
|
if len(got) != len(want) || got[0] != want[0] || got[1] != want[1] {
|
|
t.Errorf("got %#v, want %#v", got, want)
|
|
}
|
|
}
|
|
|
|
// TestGeminiAdapter_ListModels_FollowsPagination is the M7 regression
|
|
// test: a non-empty nextPageToken must make listModels fetch a second
|
|
// page via ?pageToken=<token> and return the union of both pages.
|
|
func TestGeminiAdapter_ListModels_FollowsPagination(t *testing.T) {
|
|
var gotPageTokens, gotPageSizes []string
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
gotPageTokens = append(gotPageTokens, r.URL.Query().Get("pageToken"))
|
|
gotPageSizes = append(gotPageSizes, r.URL.Query().Get("pageSize"))
|
|
w.Header().Set("Content-Type", "application/json")
|
|
if r.URL.Query().Get("pageToken") == "" {
|
|
_, _ = w.Write([]byte(`{"models":[{"name":"models/gemini-2.5-pro"},{"name":"models/gemini-2.5-flash"}],"nextPageToken":"page2"}`))
|
|
return
|
|
}
|
|
_, _ = w.Write([]byte(`{"models":[{"name":"models/gemini-2.0-flash"}]}`))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
a, err := newGeminiAdapter("p1", srv.URL, "AIza-test")
|
|
if err != nil {
|
|
t.Fatalf("newGeminiAdapter: %v", err)
|
|
}
|
|
|
|
got, err := a.listModels(context.Background())
|
|
if err != nil {
|
|
t.Fatalf("listModels: %v", err)
|
|
}
|
|
want := []string{"gemini-2.5-pro", "gemini-2.5-flash", "gemini-2.0-flash"}
|
|
if len(got) != len(want) {
|
|
t.Fatalf("got %#v, want %#v", got, want)
|
|
}
|
|
for i := range want {
|
|
if got[i] != want[i] {
|
|
t.Errorf("got[%d] = %q, want %q", i, got[i], want[i])
|
|
}
|
|
}
|
|
if len(gotPageTokens) != 2 || gotPageTokens[0] != "" || gotPageTokens[1] != "page2" {
|
|
t.Errorf("pageToken sequence = %#v, want [\"\", \"page2\"] (first page unqualified, second page's pageToken = first page's nextPageToken)", gotPageTokens)
|
|
}
|
|
// Gemini requires identical parameters on every page of one listing,
|
|
// so pageSize must ride along with the pageToken too.
|
|
for i, s := range gotPageSizes {
|
|
if s != "1000" {
|
|
t.Errorf("page %d pageSize = %q, want \"1000\"", i, s)
|
|
}
|
|
}
|
|
}
|
|
|
|
// TestGeminiAdapter_ListModels_PageCapStopsRunawayPagination proves the
|
|
// M7 fix's other half: an upstream that always returns a nextPageToken
|
|
// does not page forever — listModels stops after
|
|
// geminiListModelsPageCap requests and returns successfully with
|
|
// whatever it accumulated.
|
|
func TestGeminiAdapter_ListModels_PageCapStopsRunawayPagination(t *testing.T) {
|
|
var requests int
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
requests++
|
|
w.Header().Set("Content-Type", "application/json")
|
|
fmt.Fprintf(w, `{"models":[{"name":"models/model-%d"}],"nextPageToken":"tok-%d"}`, requests, requests)
|
|
}))
|
|
defer srv.Close()
|
|
|
|
a, err := newGeminiAdapter("p1", srv.URL, "AIza-test")
|
|
if err != nil {
|
|
t.Fatalf("newGeminiAdapter: %v", err)
|
|
}
|
|
|
|
got, err := a.listModels(context.Background())
|
|
if err != nil {
|
|
t.Fatalf("listModels: %v", err)
|
|
}
|
|
if requests != geminiListModelsPageCap {
|
|
t.Errorf("upstream received %d requests, want exactly %d (the page cap)", requests, geminiListModelsPageCap)
|
|
}
|
|
if len(got) != geminiListModelsPageCap {
|
|
t.Errorf("got %d ids, want %d (one per page up to the cap)", len(got), geminiListModelsPageCap)
|
|
}
|
|
}
|
|
|
|
// TestGeminiAdapter_HTTPClient proves httpClient returns the adapter's own
|
|
// *http.Client — the one newGeminiAdapter constructed via
|
|
// newAdapterHTTPClient — not a fresh or nil client.
|
|
func TestGeminiAdapter_HTTPClient(t *testing.T) {
|
|
a, err := newGeminiAdapter("p1", "https://generativelanguage.googleapis.com", "AIza-test")
|
|
require.NoError(t, err)
|
|
require.NotNil(t, a.httpClient())
|
|
assert.Same(t, a.client, a.httpClient())
|
|
}
|