1package provider
2
3import (
4 "context"
5 "encoding/base64"
6 "encoding/json"
7 "io"
8 "net/http"
9 "net/http/httptest"
10 "os"
11 "strings"
12 "testing"
13)
14
15// TestGeminiStreamHappyPath drives the gemini client end-to-end against
16// a fake SSE server speaking the Gemini Generative Language wire format.
17// We assert text deltas accumulate, usage rolls in from usageMetadata,
18// and the final EventDone carries StopEnd.
19func TestGeminiStreamHappyPath(t *testing.T) {
20 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
21 if got := r.Header.Get("x-goog-api-key"); got != "k" {
22 t.Errorf("missing api key header: %q", got)
23 }
24 if !strings.Contains(r.URL.RawQuery, "alt=sse") {
25 t.Errorf("missing alt=sse: %q", r.URL.RawQuery)
26 }
27 w.Header().Set("content-type", "text/event-stream")
28 w.WriteHeader(http.StatusOK)
29 fl, _ := w.(http.Flusher)
30 write := func(s string) {
31 _, _ = w.Write([]byte(s))
32 if fl != nil {
33 fl.Flush()
34 }
35 }
36 // Two text chunks, then a usage-only finish chunk.
37 write("data: " + `{"candidates":[{"content":{"role":"model","parts":[{"text":"hel"}]}}]}` + "\n\n")
38 write("data: " + `{"candidates":[{"content":{"role":"model","parts":[{"text":"lo"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":12,"candidatesTokenCount":2,"totalTokenCount":14}}` + "\n\n")
39 }))
40 defer srv.Close()
41
42 c := NewGemini("k", srv.URL)
43 evs, err := c.Stream(context.Background(), Request{Model: "gemini-2.5-pro"})
44 if err != nil {
45 t.Fatal(err)
46 }
47 var gotText string
48 var done EventDone
49 var usage Usage
50 for ev := range evs {
51 switch e := ev.(type) {
52 case EventTextDelta:
53 gotText += e.Delta
54 case EventUsage:
55 usage = e.Usage
56 case EventDone:
57 done = e
58 }
59 }
60 if gotText != "hello" {
61 t.Fatalf("text=%q", gotText)
62 }
63 if done.Stop != StopEnd {
64 t.Fatalf("stop=%v", done.Stop)
65 }
66 if usage.InputTokens != 12 || usage.OutputTokens != 2 {
67 t.Fatalf("usage=%+v", usage)
68 }
69}
70
71func TestGeminiStreamInlineImage(t *testing.T) {
72 dir := t.TempDir()
73 prev, err := os.Getwd()
74 if err != nil {
75 t.Fatal(err)
76 }
77 if err := os.Chdir(dir); err != nil {
78 t.Fatal(err)
79 }
80 t.Cleanup(func() { _ = os.Chdir(prev) })
81 img := base64.StdEncoding.EncodeToString([]byte("png-bytes"))
82 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
83 w.Header().Set("content-type", "text/event-stream")
84 w.WriteHeader(http.StatusOK)
85 _, _ = w.Write([]byte("data: " + `{"candidates":[{"content":{"role":"model","parts":[{"text":"here"},{"inlineData":{"mimeType":"image/png","data":"` + img + `"}}]},"finishReason":"STOP"}]}` + "\n\n"))
86 }))
87 defer srv.Close()
88
89 c := NewGemini("k", srv.URL)
90 evs, err := c.Stream(context.Background(), Request{Model: "gemini-2.5-flash-image"})
91 if err != nil {
92 t.Fatal(err)
93 }
94 var done EventDone
95 for ev := range evs {
96 if e, ok := ev.(EventDone); ok {
97 done = e
98 }
99 }
100 if len(done.Message.Content) != 3 {
101 t.Fatalf("content count=%d: %+v", len(done.Message.Content), done.Message.Content)
102 }
103 if tb, ok := done.Message.Content[0].(TextBlock); !ok || tb.Text != "here" {
104 t.Fatalf("text block=%T %+v", done.Message.Content[0], done.Message.Content[0])
105 }
106 ib, ok := done.Message.Content[1].(ImageBlock)
107 if !ok {
108 t.Fatalf("want image block, got %T", done.Message.Content[1])
109 }
110 if ib.MimeType != "image/png" || string(ib.Data) != "png-bytes" {
111 t.Fatalf("image=%q %q", ib.MimeType, string(ib.Data))
112 }
113 saved, ok := done.Message.Content[2].(TextBlock)
114 if !ok || !strings.Contains(saved.Text, "zot-gemini-image-") || !strings.Contains(saved.Text, ".png") {
115 t.Fatalf("saved path block=%T %+v", done.Message.Content[2], done.Message.Content[2])
116 }
117 path := strings.TrimPrefix(saved.Text, "Saved image: `")
118 path = strings.TrimSuffix(path, "`")
119 if _, err := os.Stat(path); err != nil {
120 t.Fatalf("saved image missing at %q: %v", path, err)
121 }
122}
123
124// TestGeminiToolCall covers the tool-call branch: a single
125// functionCall part should produce ToolStart/Args/End and the final
126// stop reason should be StopToolUse.
127func TestGeminiToolCall(t *testing.T) {
128 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
129 w.Header().Set("content-type", "text/event-stream")
130 w.WriteHeader(http.StatusOK)
131 fl, _ := w.(http.Flusher)
132 write := func(s string) {
133 _, _ = w.Write([]byte(s))
134 if fl != nil {
135 fl.Flush()
136 }
137 }
138 write("data: " + `{"candidates":[{"content":{"role":"model","parts":[{"functionCall":{"name":"read","args":{"path":"a"}}}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":3,"candidatesTokenCount":7,"totalTokenCount":10}}` + "\n\n")
139 }))
140 defer srv.Close()
141
142 c := NewGemini("k", srv.URL)
143 evs, err := c.Stream(context.Background(), Request{Model: "gemini-2.5-pro"})
144 if err != nil {
145 t.Fatal(err)
146 }
147 var startName, argsBuf string
148 var done EventDone
149 for ev := range evs {
150 switch e := ev.(type) {
151 case EventToolStart:
152 startName = e.Name
153 case EventToolArgs:
154 argsBuf += e.Delta
155 case EventDone:
156 done = e
157 }
158 }
159 if startName != "read" {
160 t.Fatalf("tool name=%q", startName)
161 }
162 if !strings.Contains(argsBuf, `"path":"a"`) {
163 t.Fatalf("tool args=%q", argsBuf)
164 }
165 if done.Stop != StopToolUse {
166 t.Fatalf("stop=%v", done.Stop)
167 }
168 // The assembled message should contain the tool call as the first content block.
169 if len(done.Message.Content) != 1 {
170 t.Fatalf("message content count=%d", len(done.Message.Content))
171 }
172 tc, ok := done.Message.Content[0].(ToolCallBlock)
173 if !ok {
174 t.Fatalf("expected ToolCallBlock, got %T", done.Message.Content[0])
175 }
176 if tc.Name != "read" {
177 t.Fatalf("tool block name=%q", tc.Name)
178 }
179}
180
181// TestGeminiBuildRequestSystemAndTools confirms the wire payload puts
182// the system prompt under systemInstruction and tool defs under tools[0].
183func TestGeminiBuildRequestSystemAndTools(t *testing.T) {
184 c := NewGemini("k", "https://example.invalid").(*geminiClient)
185 wire, _, err := c.buildRequest(Request{
186 Model: "gemini-2.5-pro",
187 System: "you are zot",
188 Tools: []Tool{
189 {Name: "read", Description: "read a file", Schema: json.RawMessage(`{"type":"object","properties":{"path":{"type":"string"}}}`)},
190 },
191 Messages: []Message{{Role: RoleUser, Content: []Content{TextBlock{Text: "hi"}}}},
192 })
193 if err != nil {
194 t.Fatal(err)
195 }
196 if wire.SystemInstruction == nil || wire.SystemInstruction.Parts[0].Text != "you are zot" {
197 t.Fatalf("system: %+v", wire.SystemInstruction)
198 }
199 if len(wire.Tools) != 1 || len(wire.Tools[0].FunctionDeclarations) != 1 || wire.Tools[0].FunctionDeclarations[0].Name != "read" {
200 t.Fatalf("tools: %+v", wire.Tools)
201 }
202 if len(wire.Contents) != 1 || wire.Contents[0].Role != "user" {
203 t.Fatalf("contents: %+v", wire.Contents)
204 }
205}
206
207// TestGeminiBuildRequestImageModelOmitsTools confirms image-generation
208// models receive direct multimodal prompts without function declarations.
209func TestGeminiBuildRequestStripsUnsupportedSchemaFields(t *testing.T) {
210 c := NewGemini("k", "https://example.invalid").(*geminiClient)
211 wire, _, err := c.buildRequest(Request{
212 Model: "gemini-2.5-pro",
213 Tools: []Tool{{
214 Name: "edit",
215 Description: "edit a file",
216 Schema: json.RawMessage(`{
217 "$schema":"http://json-schema.org/draft-07/schema#",
218 "type":"object",
219 "additionalProperties":false,
220 "properties":{
221 "edits":{
222 "type":"array",
223 "items":{
224 "type":"object",
225 "additionalProperties":false,
226 "properties":{"oldText":{"type":"string"},"newText":{"type":"string"}}
227 }
228 }
229 }
230 }`),
231 }},
232 })
233 if err != nil {
234 t.Fatal(err)
235 }
236 got := string(wire.Tools[0].FunctionDeclarations[0].Parameters)
237 if strings.Contains(got, "additionalProperties") || strings.Contains(got, "$schema") {
238 t.Fatalf("Gemini schema should strip unsupported fields, got %s", got)
239 }
240 if !strings.Contains(got, `"oldText"`) || !strings.Contains(got, `"newText"`) {
241 t.Fatalf("Gemini schema lost nested properties, got %s", got)
242 }
243}
244
245func TestGeminiBuildRequestImageModelOmitsTools(t *testing.T) {
246 c := NewGemini("k", "https://example.invalid").(*geminiClient)
247 wire, _, err := c.buildRequest(Request{
248 Model: "gemini-2.5-flash-image",
249 Tools: []Tool{
250 {Name: "read", Description: "read a file", Schema: json.RawMessage(`{"type":"object","properties":{"path":{"type":"string"}}}`)},
251 },
252 Messages: []Message{
253 {Role: RoleUser, Content: []Content{TextBlock{Text: "edit this image"}, ImageBlock{MimeType: "image/png", Data: []byte("png")}}},
254 {Role: RoleAssistant, Content: []Content{TextBlock{Text: "checking"}, ToolCallBlock{ID: "1", Name: "read", Arguments: json.RawMessage(`{"path":"x"}`)}}},
255 {Role: RoleTool, Content: []Content{ToolResultBlock{CallID: "1", Content: []Content{TextBlock{Text: "tool output"}}}}},
256 },
257 })
258 if err != nil {
259 t.Fatal(err)
260 }
261 if len(wire.Tools) != 0 {
262 t.Fatalf("image model should omit tools, got %+v", wire.Tools)
263 }
264 if len(wire.Contents) != 2 {
265 t.Fatalf("contents count=%d: %+v", len(wire.Contents), wire.Contents)
266 }
267 if len(wire.Contents[0].Parts) != 2 || wire.Contents[0].Parts[1].InlineData == nil {
268 t.Fatalf("user image parts not preserved: %+v", wire.Contents[0].Parts)
269 }
270 if len(wire.Contents[1].Parts) != 1 || wire.Contents[1].Parts[0].FunctionCall != nil {
271 t.Fatalf("assistant tool call not stripped: %+v", wire.Contents[1].Parts)
272 }
273}
274
275// TestGeminiErrorStatus confirms HTTP error bodies bubble up.
276func TestGeminiErrorStatus(t *testing.T) {
277 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
278 w.WriteHeader(http.StatusUnauthorized)
279 _, _ = io.WriteString(w, `{"error":{"code":401,"message":"bad key"}}`)
280 }))
281 defer srv.Close()
282
283 c := NewGemini("k", srv.URL)
284 _, err := c.Stream(context.Background(), Request{Model: "gemini-2.5-pro"})
285 if err == nil || !strings.Contains(err.Error(), "401") {
286 t.Fatalf("want 401 err, got %v", err)
287 }
288}
289
290// TestGeminiThinkingConfig spot-checks the level → wire mapping for
291// representative model families. The full table is exercised via
292// integration; this guards the routing logic.
293func TestGeminiThinkingConfig(t *testing.T) {
294 cases := []struct {
295 modelID string
296 level string
297 wantLvl string
298 wantBud int
299 }{
300 {"gemini-3-pro", "low", "LOW", 0},
301 {"gemini-3-pro", "medium", "HIGH", 0}, // Pro can't go below LOW; medium → HIGH
302 {"gemini-3-flash", "medium", "MEDIUM", 0},
303 {"gemini-2.5-pro", "high", "", 16384},
304 {"gemini-2.5-pro", "maximum", "", 32768},
305 {"gemini-2.5-flash", "minimum", "", 1024},
306 {"gemini-2.5-flash", "low", "", 2048},
307 {"gemini-2.0-flash", "high", "", 0}, // 2.0 has no thinking → nil
308 }
309 for _, tc := range cases {
310 got := geminiThinkingConfig(tc.modelID, tc.level)
311 if tc.wantLvl == "" && tc.wantBud == 0 {
312 if got != nil {
313 t.Errorf("%s/%s: want nil, got %+v", tc.modelID, tc.level, got)
314 }
315 continue
316 }
317 if got == nil {
318 t.Errorf("%s/%s: got nil", tc.modelID, tc.level)
319 continue
320 }
321 if tc.wantLvl != "" && got.ThinkingLevel != tc.wantLvl {
322 t.Errorf("%s/%s: level=%q want %q", tc.modelID, tc.level, got.ThinkingLevel, tc.wantLvl)
323 }
324 if tc.wantBud != 0 {
325 if got.ThinkingBudget == nil || *got.ThinkingBudget != tc.wantBud {
326 t.Errorf("%s/%s: budget=%v want %d", tc.modelID, tc.level, got.ThinkingBudget, tc.wantBud)
327 }
328 }
329 }
330}
331
332// TestDiscoverGoogle exercises the discovery helper against a fake
333// /v1beta/models endpoint, confirming pagination plus filtering of
334// non-chat ids (embedding, aqa).
335func TestDiscoverGoogle(t *testing.T) {
336 calls := 0
337 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
338 if r.Header.Get("x-goog-api-key") != "k" {
339 t.Errorf("missing api key")
340 }
341 calls++
342 if calls == 1 {
343 _, _ = io.WriteString(w, `{
344 "models": [
345 {"name":"models/gemini-2.5-pro","displayName":"Gemini 2.5 Pro","supportedGenerationMethods":["generateContent","streamGenerateContent"],"inputTokenLimit":1048576,"outputTokenLimit":65536},
346 {"name":"models/text-embedding-004","displayName":"Text Embedding","supportedGenerationMethods":["embedContent"]},
347 {"name":"models/aqa","displayName":"AQA","supportedGenerationMethods":["generateAnswer"]}
348 ],
349 "nextPageToken": "p2"
350 }`)
351 } else {
352 _, _ = io.WriteString(w, `{
353 "models": [
354 {"name":"models/gemini-2.5-flash","displayName":"Gemini 2.5 Flash","supportedGenerationMethods":["generateContent"]}
355 ]
356 }`)
357 }
358 }))
359 defer srv.Close()
360
361 got, err := DiscoverGoogle(context.Background(), "k", srv.URL)
362 if err != nil {
363 t.Fatal(err)
364 }
365 if calls != 2 {
366 t.Errorf("want 2 calls, got %d", calls)
367 }
368 if len(got) != 2 {
369 t.Fatalf("want 2 models, got %d: %+v", len(got), got)
370 }
371 if got[0].ID != "gemini-2.5-pro" || got[0].Provider != "google" {
372 t.Errorf("first model wrong: %+v", got[0])
373 }
374 if got[1].ID != "gemini-2.5-flash" {
375 t.Errorf("second model wrong: %+v", got[1])
376 }
377}