1package provider
2
3import (
4 "context"
5 "encoding/json"
6 "io"
7 "net/http"
8 "net/http/httptest"
9 "strings"
10 "testing"
11)
12
13// TestGeminiStreamHappyPath drives the gemini client end-to-end against
14// a fake SSE server speaking the Gemini Generative Language wire format.
15// We assert text deltas accumulate, usage rolls in from usageMetadata,
16// and the final EventDone carries StopEnd.
17func TestGeminiStreamHappyPath(t *testing.T) {
18 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
19 if got := r.Header.Get("x-goog-api-key"); got != "k" {
20 t.Errorf("missing api key header: %q", got)
21 }
22 if !strings.Contains(r.URL.RawQuery, "alt=sse") {
23 t.Errorf("missing alt=sse: %q", r.URL.RawQuery)
24 }
25 w.Header().Set("content-type", "text/event-stream")
26 w.WriteHeader(http.StatusOK)
27 fl, _ := w.(http.Flusher)
28 write := func(s string) {
29 _, _ = w.Write([]byte(s))
30 if fl != nil {
31 fl.Flush()
32 }
33 }
34 // Two text chunks, then a usage-only finish chunk.
35 write("data: " + `{"candidates":[{"content":{"role":"model","parts":[{"text":"hel"}]}}]}` + "\n\n")
36 write("data: " + `{"candidates":[{"content":{"role":"model","parts":[{"text":"lo"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":12,"candidatesTokenCount":2,"totalTokenCount":14}}` + "\n\n")
37 }))
38 defer srv.Close()
39
40 c := NewGemini("k", srv.URL)
41 evs, err := c.Stream(context.Background(), Request{Model: "gemini-2.5-pro"})
42 if err != nil {
43 t.Fatal(err)
44 }
45 var gotText string
46 var done EventDone
47 var usage Usage
48 for ev := range evs {
49 switch e := ev.(type) {
50 case EventTextDelta:
51 gotText += e.Delta
52 case EventUsage:
53 usage = e.Usage
54 case EventDone:
55 done = e
56 }
57 }
58 if gotText != "hello" {
59 t.Fatalf("text=%q", gotText)
60 }
61 if done.Stop != StopEnd {
62 t.Fatalf("stop=%v", done.Stop)
63 }
64 if usage.InputTokens != 12 || usage.OutputTokens != 2 {
65 t.Fatalf("usage=%+v", usage)
66 }
67}
68
69// TestGeminiToolCall covers the tool-call branch: a single
70// functionCall part should produce ToolStart/Args/End and the final
71// stop reason should be StopToolUse.
72func TestGeminiToolCall(t *testing.T) {
73 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
74 w.Header().Set("content-type", "text/event-stream")
75 w.WriteHeader(http.StatusOK)
76 fl, _ := w.(http.Flusher)
77 write := func(s string) {
78 _, _ = w.Write([]byte(s))
79 if fl != nil {
80 fl.Flush()
81 }
82 }
83 write("data: " + `{"candidates":[{"content":{"role":"model","parts":[{"functionCall":{"name":"read","args":{"path":"a"}}}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":3,"candidatesTokenCount":7,"totalTokenCount":10}}` + "\n\n")
84 }))
85 defer srv.Close()
86
87 c := NewGemini("k", srv.URL)
88 evs, err := c.Stream(context.Background(), Request{Model: "gemini-2.5-pro"})
89 if err != nil {
90 t.Fatal(err)
91 }
92 var startName, argsBuf string
93 var done EventDone
94 for ev := range evs {
95 switch e := ev.(type) {
96 case EventToolStart:
97 startName = e.Name
98 case EventToolArgs:
99 argsBuf += e.Delta
100 case EventDone:
101 done = e
102 }
103 }
104 if startName != "read" {
105 t.Fatalf("tool name=%q", startName)
106 }
107 if !strings.Contains(argsBuf, `"path":"a"`) {
108 t.Fatalf("tool args=%q", argsBuf)
109 }
110 if done.Stop != StopToolUse {
111 t.Fatalf("stop=%v", done.Stop)
112 }
113 // The assembled message should contain the tool call as the first content block.
114 if len(done.Message.Content) != 1 {
115 t.Fatalf("message content count=%d", len(done.Message.Content))
116 }
117 tc, ok := done.Message.Content[0].(ToolCallBlock)
118 if !ok {
119 t.Fatalf("expected ToolCallBlock, got %T", done.Message.Content[0])
120 }
121 if tc.Name != "read" {
122 t.Fatalf("tool block name=%q", tc.Name)
123 }
124}
125
126// TestGeminiBuildRequestSystemAndTools confirms the wire payload puts
127// the system prompt under systemInstruction and tool defs under tools[0].
128func TestGeminiBuildRequestSystemAndTools(t *testing.T) {
129 c := NewGemini("k", "https://example.invalid").(*geminiClient)
130 wire, _, err := c.buildRequest(Request{
131 Model: "gemini-2.5-pro",
132 System: "you are zot",
133 Tools: []Tool{
134 {Name: "read", Description: "read a file", Schema: json.RawMessage(`{"type":"object","properties":{"path":{"type":"string"}}}`)},
135 },
136 Messages: []Message{{Role: RoleUser, Content: []Content{TextBlock{Text: "hi"}}}},
137 })
138 if err != nil {
139 t.Fatal(err)
140 }
141 if wire.SystemInstruction == nil || wire.SystemInstruction.Parts[0].Text != "you are zot" {
142 t.Fatalf("system: %+v", wire.SystemInstruction)
143 }
144 if len(wire.Tools) != 1 || len(wire.Tools[0].FunctionDeclarations) != 1 || wire.Tools[0].FunctionDeclarations[0].Name != "read" {
145 t.Fatalf("tools: %+v", wire.Tools)
146 }
147 if len(wire.Contents) != 1 || wire.Contents[0].Role != "user" {
148 t.Fatalf("contents: %+v", wire.Contents)
149 }
150}
151
152// TestGeminiErrorStatus confirms HTTP error bodies bubble up.
153func TestGeminiErrorStatus(t *testing.T) {
154 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
155 w.WriteHeader(http.StatusUnauthorized)
156 _, _ = io.WriteString(w, `{"error":{"code":401,"message":"bad key"}}`)
157 }))
158 defer srv.Close()
159
160 c := NewGemini("k", srv.URL)
161 _, err := c.Stream(context.Background(), Request{Model: "gemini-2.5-pro"})
162 if err == nil || !strings.Contains(err.Error(), "401") {
163 t.Fatalf("want 401 err, got %v", err)
164 }
165}
166
167// TestGeminiThinkingConfig spot-checks the level → wire mapping for
168// representative model families. The full table is exercised via
169// integration; this guards the routing logic.
170func TestGeminiThinkingConfig(t *testing.T) {
171 cases := []struct {
172 modelID string
173 level string
174 wantLvl string
175 wantBud int
176 }{
177 {"gemini-3-pro", "low", "LOW", 0},
178 {"gemini-3-pro", "medium", "HIGH", 0}, // Pro can't go below LOW; medium → HIGH
179 {"gemini-3-flash", "medium", "MEDIUM", 0},
180 {"gemini-2.5-pro", "high", "", 32768},
181 {"gemini-2.5-flash", "low", "", 2048},
182 {"gemini-2.0-flash", "high", "", 0}, // 2.0 has no thinking → nil
183 }
184 for _, tc := range cases {
185 got := geminiThinkingConfig(tc.modelID, tc.level)
186 if tc.wantLvl == "" && tc.wantBud == 0 {
187 if got != nil {
188 t.Errorf("%s/%s: want nil, got %+v", tc.modelID, tc.level, got)
189 }
190 continue
191 }
192 if got == nil {
193 t.Errorf("%s/%s: got nil", tc.modelID, tc.level)
194 continue
195 }
196 if tc.wantLvl != "" && got.ThinkingLevel != tc.wantLvl {
197 t.Errorf("%s/%s: level=%q want %q", tc.modelID, tc.level, got.ThinkingLevel, tc.wantLvl)
198 }
199 if tc.wantBud != 0 {
200 if got.ThinkingBudget == nil || *got.ThinkingBudget != tc.wantBud {
201 t.Errorf("%s/%s: budget=%v want %d", tc.modelID, tc.level, got.ThinkingBudget, tc.wantBud)
202 }
203 }
204 }
205}
206
207// TestDiscoverGoogle exercises the discovery helper against a fake
208// /v1beta/models endpoint, confirming pagination plus filtering of
209// non-chat ids (embedding, aqa).
210func TestDiscoverGoogle(t *testing.T) {
211 calls := 0
212 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
213 if r.Header.Get("x-goog-api-key") != "k" {
214 t.Errorf("missing api key")
215 }
216 calls++
217 if calls == 1 {
218 _, _ = io.WriteString(w, `{
219 "models": [
220 {"name":"models/gemini-2.5-pro","displayName":"Gemini 2.5 Pro","supportedGenerationMethods":["generateContent","streamGenerateContent"],"inputTokenLimit":1048576,"outputTokenLimit":65536},
221 {"name":"models/text-embedding-004","displayName":"Text Embedding","supportedGenerationMethods":["embedContent"]},
222 {"name":"models/aqa","displayName":"AQA","supportedGenerationMethods":["generateAnswer"]}
223 ],
224 "nextPageToken": "p2"
225 }`)
226 } else {
227 _, _ = io.WriteString(w, `{
228 "models": [
229 {"name":"models/gemini-2.5-flash","displayName":"Gemini 2.5 Flash","supportedGenerationMethods":["generateContent"]}
230 ]
231 }`)
232 }
233 }))
234 defer srv.Close()
235
236 got, err := DiscoverGoogle(context.Background(), "k", srv.URL)
237 if err != nil {
238 t.Fatal(err)
239 }
240 if calls != 2 {
241 t.Errorf("want 2 calls, got %d", calls)
242 }
243 if len(got) != 2 {
244 t.Fatalf("want 2 models, got %d: %+v", len(got), got)
245 }
246 if got[0].ID != "gemini-2.5-pro" || got[0].Provider != "google" {
247 t.Errorf("first model wrong: %+v", got[0])
248 }
249 if got[1].ID != "gemini-2.5-flash" {
250 t.Errorf("second model wrong: %+v", got[1])
251 }
252}