package gemini import ( "context" "encoding/json" "io" "net/http" "net/http/httptest" "strings" "testing" "github.com/zzet/gortex/internal/llm" ) func TestNew_MissingKey(t *testing.T) { t.Setenv("GEMINI_API_KEY", "") if _, err := New(llm.RemoteConfig{Model: "gemini-2.5-pro"}); err == nil { t.Fatal("expected error when API key env is unset") } } func TestNew_MissingModel(t *testing.T) { t.Setenv("GEMINI_API_KEY", "k") if _, err := New(llm.RemoteConfig{}); err == nil { t.Fatal("expected error when model is unset") } } func TestComplete_StructuredUsesResponseSchema(t *testing.T) { var gotBody map[string]any var gotPath string srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { gotPath = r.URL.Path if r.Header.Get("x-goog-api-key") != "test-key" { t.Errorf("x-goog-api-key=%q", r.Header.Get("x-goog-api-key")) } body, _ := io.ReadAll(r.Body) _ = json.Unmarshal(body, &gotBody) _, _ = io.WriteString(w, `{"candidates":[{"content":{"parts":[{"text":"{\"terms\":[\"jwt\"]}"}],"role":"model"}}]}`) })) defer srv.Close() t.Setenv("GEMINI_API_KEY", "test-key") p, err := New(llm.RemoteConfig{Model: "gemini-2.5-pro", BaseURL: srv.URL}) if err != nil { t.Fatal(err) } defer p.Close() resp, err := p.Complete(context.Background(), llm.CompletionRequest{ Messages: []llm.Message{ {Role: llm.RoleSystem, Content: "you expand queries"}, {Role: llm.RoleUser, Content: "Query: auth"}, }, Shape: llm.ShapeExpandTerms, MaxTokens: 64, }) if err != nil { t.Fatal(err) } if resp.Text != `{"terms":["jwt"]}` { t.Errorf("text=%q", resp.Text) } if !strings.Contains(gotPath, "gemini-2.5-pro:generateContent") { t.Errorf("path=%q does not embed model + action", gotPath) } sys, _ := gotBody["system_instruction"].(map[string]any) if sys == nil { t.Fatal("system_instruction missing") } gen, _ := gotBody["generationConfig"].(map[string]any) if gen == nil { t.Fatal("generationConfig missing") } if gen["responseMimeType"] != "application/json" { t.Errorf("responseMimeType=%v want application/json", gen["responseMimeType"]) } schema, _ := gen["responseSchema"].(map[string]any) if schema == nil { t.Fatal("responseSchema missing on structured request") } if _, ok := schema["additionalProperties"]; ok { t.Error("responseSchema must not carry additionalProperties — Gemini rejects it") } } func TestComplete_FreeformNoSchema(t *testing.T) { var gotBody map[string]any srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { body, _ := io.ReadAll(r.Body) _ = json.Unmarshal(body, &gotBody) _, _ = io.WriteString(w, `{"candidates":[{"content":{"parts":[{"text":"hi"}],"role":"model"}}]}`) })) defer srv.Close() t.Setenv("GEMINI_API_KEY", "k") p, _ := New(llm.RemoteConfig{Model: "m", BaseURL: srv.URL}) defer p.Close() resp, err := p.Complete(context.Background(), llm.CompletionRequest{ Messages: []llm.Message{{Role: llm.RoleUser, Content: "hi"}}, Shape: llm.ShapeFreeform, }) if err != nil { t.Fatal(err) } if resp.Text != "hi" { t.Errorf("text=%q", resp.Text) } if gen, ok := gotBody["generationConfig"].(map[string]any); ok { if gen["responseSchema"] != nil { t.Errorf("freeform request must not set responseSchema, got %v", gen["responseSchema"]) } } } func TestComplete_ToolResultBecomesUserTurn(t *testing.T) { var gotBody map[string]any srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { body, _ := io.ReadAll(r.Body) _ = json.Unmarshal(body, &gotBody) _, _ = io.WriteString(w, `{"candidates":[{"content":{"parts":[{"text":"ok"}]}}]}`) })) defer srv.Close() t.Setenv("GEMINI_API_KEY", "k") p, _ := New(llm.RemoteConfig{Model: "m", BaseURL: srv.URL}) defer p.Close() _, err := p.Complete(context.Background(), llm.CompletionRequest{ Messages: []llm.Message{ {Role: llm.RoleUser, Content: "q"}, {Role: llm.RoleAssistant, Content: `{"tool":"x"}`}, {Role: llm.RoleTool, Content: `{"result":1}`, ToolName: "x"}, }, Shape: llm.ShapeFreeform, }) if err != nil { t.Fatal(err) } contents := gotBody["contents"].([]any) if len(contents) != 3 { t.Fatalf("contents=%d want 3", len(contents)) } asst := contents[1].(map[string]any) if asst["role"] != "model" { t.Errorf("assistant role=%v want model (Gemini uses 'model' for assistant turns)", asst["role"]) } tool := contents[2].(map[string]any) if tool["role"] != "user" { t.Errorf("tool result role=%v want user (emulated protocol)", tool["role"]) } } func TestComplete_APIError(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusBadRequest) _, _ = io.WriteString(w, `{"error":{"code":400,"message":"bad model","status":"INVALID_ARGUMENT"}}`) })) defer srv.Close() t.Setenv("GEMINI_API_KEY", "k") p, _ := New(llm.RemoteConfig{Model: "m", BaseURL: srv.URL}) defer p.Close() if _, err := p.Complete(context.Background(), llm.CompletionRequest{ Messages: []llm.Message{{Role: llm.RoleUser, Content: "hi"}}, }); err == nil { t.Fatal("expected an error for a non-200 response") } } func TestName(t *testing.T) { t.Setenv("GEMINI_API_KEY", "k") p, _ := New(llm.RemoteConfig{Model: "m"}) if p.Name() != "gemini" { t.Errorf("Name()=%q", p.Name()) } } // TestComplete_HollowResponseRetriesThenSucceeds proves an HTTP 200 // with an empty candidates array is treated as a transient truncation: // the first hollow answer is retried and the second, good answer wins. func TestComplete_HollowResponseRetriesThenSucceeds(t *testing.T) { var calls int srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { calls++ if calls == 1 { _, _ = io.WriteString(w, `{"candidates":[]}`) return } _, _ = io.WriteString(w, `{"candidates":[{"content":{"parts":[{"text":"recovered"}]}}]}`) })) defer srv.Close() t.Setenv("GEMINI_API_KEY", "k") p, _ := New(llm.RemoteConfig{Model: "m", BaseURL: srv.URL}) defer p.Close() resp, err := p.Complete(context.Background(), llm.CompletionRequest{ Messages: []llm.Message{{Role: llm.RoleUser, Content: "hi"}}, }) if err != nil { t.Fatalf("unexpected error: %v", err) } if resp.Text != "recovered" { t.Errorf("text=%q want %q", resp.Text, "recovered") } if calls != 2 { t.Errorf("calls=%d want 2 (one hollow 200, then a retry)", calls) } } func TestSanitiseSchema_StripsAdditionalProperties(t *testing.T) { in := map[string]any{ "type": "object", "properties": map[string]any{ "terms": map[string]any{ "type": "array", "items": map[string]any{"type": "string"}, "additionalProperties": false, }, }, "additionalProperties": false, } out := sanitiseSchema(in) if _, ok := out["additionalProperties"]; ok { t.Error("top-level additionalProperties not stripped") } props := out["properties"].(map[string]any) terms := props["terms"].(map[string]any) if _, ok := terms["additionalProperties"]; ok { t.Error("nested additionalProperties not stripped") } }