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 TestGeminiBuildRequestImageModelOmitsTools(t *testing.T) {
210 c := NewGemini("k", "https://example.invalid").(*geminiClient)
211 wire, _, err := c.buildRequest(Request{
212 Model: "gemini-2.5-flash-image",
213 Tools: []Tool{
214 {Name: "read", Description: "read a file", Schema: json.RawMessage(`{"type":"object","properties":{"path":{"type":"string"}}}`)},
215 },
216 Messages: []Message{
217 {Role: RoleUser, Content: []Content{TextBlock{Text: "edit this image"}, ImageBlock{MimeType: "image/png", Data: []byte("png")}}},
218 {Role: RoleAssistant, Content: []Content{TextBlock{Text: "checking"}, ToolCallBlock{ID: "1", Name: "read", Arguments: json.RawMessage(`{"path":"x"}`)}}},
219 {Role: RoleTool, Content: []Content{ToolResultBlock{CallID: "1", Content: []Content{TextBlock{Text: "tool output"}}}}},
220 },
221 })
222 if err != nil {
223 t.Fatal(err)
224 }
225 if len(wire.Tools) != 0 {
226 t.Fatalf("image model should omit tools, got %+v", wire.Tools)
227 }
228 if len(wire.Contents) != 2 {
229 t.Fatalf("contents count=%d: %+v", len(wire.Contents), wire.Contents)
230 }
231 if len(wire.Contents[0].Parts) != 2 || wire.Contents[0].Parts[1].InlineData == nil {
232 t.Fatalf("user image parts not preserved: %+v", wire.Contents[0].Parts)
233 }
234 if len(wire.Contents[1].Parts) != 1 || wire.Contents[1].Parts[0].FunctionCall != nil {
235 t.Fatalf("assistant tool call not stripped: %+v", wire.Contents[1].Parts)
236 }
237}
238
239// TestGeminiErrorStatus confirms HTTP error bodies bubble up.
240func TestGeminiErrorStatus(t *testing.T) {
241 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
242 w.WriteHeader(http.StatusUnauthorized)
243 _, _ = io.WriteString(w, `{"error":{"code":401,"message":"bad key"}}`)
244 }))
245 defer srv.Close()
246
247 c := NewGemini("k", srv.URL)
248 _, err := c.Stream(context.Background(), Request{Model: "gemini-2.5-pro"})
249 if err == nil || !strings.Contains(err.Error(), "401") {
250 t.Fatalf("want 401 err, got %v", err)
251 }
252}
253
254// TestGeminiThinkingConfig spot-checks the level → wire mapping for
255// representative model families. The full table is exercised via
256// integration; this guards the routing logic.
257func TestGeminiThinkingConfig(t *testing.T) {
258 cases := []struct {
259 modelID string
260 level string
261 wantLvl string
262 wantBud int
263 }{
264 {"gemini-3-pro", "low", "LOW", 0},
265 {"gemini-3-pro", "medium", "HIGH", 0}, // Pro can't go below LOW; medium → HIGH
266 {"gemini-3-flash", "medium", "MEDIUM", 0},
267 {"gemini-2.5-pro", "high", "", 32768},
268 {"gemini-2.5-flash", "low", "", 2048},
269 {"gemini-2.0-flash", "high", "", 0}, // 2.0 has no thinking → nil
270 }
271 for _, tc := range cases {
272 got := geminiThinkingConfig(tc.modelID, tc.level)
273 if tc.wantLvl == "" && tc.wantBud == 0 {
274 if got != nil {
275 t.Errorf("%s/%s: want nil, got %+v", tc.modelID, tc.level, got)
276 }
277 continue
278 }
279 if got == nil {
280 t.Errorf("%s/%s: got nil", tc.modelID, tc.level)
281 continue
282 }
283 if tc.wantLvl != "" && got.ThinkingLevel != tc.wantLvl {
284 t.Errorf("%s/%s: level=%q want %q", tc.modelID, tc.level, got.ThinkingLevel, tc.wantLvl)
285 }
286 if tc.wantBud != 0 {
287 if got.ThinkingBudget == nil || *got.ThinkingBudget != tc.wantBud {
288 t.Errorf("%s/%s: budget=%v want %d", tc.modelID, tc.level, got.ThinkingBudget, tc.wantBud)
289 }
290 }
291 }
292}
293
294// TestDiscoverGoogle exercises the discovery helper against a fake
295// /v1beta/models endpoint, confirming pagination plus filtering of
296// non-chat ids (embedding, aqa).
297func TestDiscoverGoogle(t *testing.T) {
298 calls := 0
299 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
300 if r.Header.Get("x-goog-api-key") != "k" {
301 t.Errorf("missing api key")
302 }
303 calls++
304 if calls == 1 {
305 _, _ = io.WriteString(w, `{
306 "models": [
307 {"name":"models/gemini-2.5-pro","displayName":"Gemini 2.5 Pro","supportedGenerationMethods":["generateContent","streamGenerateContent"],"inputTokenLimit":1048576,"outputTokenLimit":65536},
308 {"name":"models/text-embedding-004","displayName":"Text Embedding","supportedGenerationMethods":["embedContent"]},
309 {"name":"models/aqa","displayName":"AQA","supportedGenerationMethods":["generateAnswer"]}
310 ],
311 "nextPageToken": "p2"
312 }`)
313 } else {
314 _, _ = io.WriteString(w, `{
315 "models": [
316 {"name":"models/gemini-2.5-flash","displayName":"Gemini 2.5 Flash","supportedGenerationMethods":["generateContent"]}
317 ]
318 }`)
319 }
320 }))
321 defer srv.Close()
322
323 got, err := DiscoverGoogle(context.Background(), "k", srv.URL)
324 if err != nil {
325 t.Fatal(err)
326 }
327 if calls != 2 {
328 t.Errorf("want 2 calls, got %d", calls)
329 }
330 if len(got) != 2 {
331 t.Fatalf("want 2 models, got %d: %+v", len(got), got)
332 }
333 if got[0].ID != "gemini-2.5-pro" || got[0].Provider != "google" {
334 t.Errorf("first model wrong: %+v", got[0])
335 }
336 if got[1].ID != "gemini-2.5-flash" {
337 t.Errorf("second model wrong: %+v", got[1])
338 }
339}