path: root/internal/provider/google/google_test.go

package google

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/thepeterstone/claudomator/internal/provider"
	"github.com/thepeterstone/claudomator/internal/retry"
)

// fakeGoogleServer replies to POST /v1beta/models/{model}:generateContent
// with the given status and raw JSON body, capturing the last decoded
// request for assertions. Modeled on
// internal/provider/anthropic/anthropic_test.go's fakeAnthropicServer,
// adapted to the Gemini generateContent wire shape (model in the URL path,
// key as a query param rather than a header).
func fakeGoogleServer(t *testing.T, status int, body string, capture *wireRequest) *httptest.Server {
	t.Helper()
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if !strings.Contains(r.URL.Path, ":generateContent") {
			t.Errorf("unexpected path %q", r.URL.Path)
		}
		if !strings.HasPrefix(r.URL.Path, "/v1beta/models/") {
			t.Errorf("expected path to start with /v1beta/models/, got %q", r.URL.Path)
		}
		if r.URL.Query().Get("key") != "test-key" {
			t.Errorf("missing/wrong key query param: %q", r.URL.Query().Get("key"))
		}
		if r.Header.Get("content-type") != "application/json" {
			t.Errorf("missing/wrong content-type header: %q", r.Header.Get("content-type"))
		}
		if capture != nil {
			if err := json.NewDecoder(r.Body).Decode(capture); err != nil {
				t.Fatalf("decode request body: %v", err)
			}
		}
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(status)
		w.Write([]byte(body))
	}))
}

func newTestProvider(url string) *Provider {
	return New("test-key", url, 5*time.Second)
}

// TestChat_PlainTextResponse covers a plain model text reply: no function
// calls, finishReason "STOP", usage translated to provider.Usage with a
// non-zero CostUSD computed from the pricing table.
func TestChat_PlainTextResponse(t *testing.T) {
	var got wireRequest
	srv := fakeGoogleServer(t, http.StatusOK, `{
		"candidates": [{
			"content": {"role": "model", "parts": [{"text": "Hello there!"}]},
			"finishReason": "STOP"
		}],
		"usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 5, "totalTokenCount": 15}
	}`, &got)
	defer srv.Close()

	p := newTestProvider(srv.URL)
	resp, err := p.Chat(context.Background(), provider.ChatRequest{
		Model:  "gemini-2.5-flash",
		System: "You are helpful.",
		Messages: []provider.Message{
			{Role: "user", Text: "hi"},
		},
	})
	if err != nil {
		t.Fatalf("Chat: %v", err)
	}
	if resp.Text != "Hello there!" {
		t.Errorf("Text: want %q got %q", "Hello there!", resp.Text)
	}
	if len(resp.ToolCalls) != 0 {
		t.Errorf("expected no tool calls, got %+v", resp.ToolCalls)
	}
	if resp.StopReason != "STOP" {
		t.Errorf("StopReason: want STOP got %q", resp.StopReason)
	}
	if resp.Usage.InputTokens != 10 || resp.Usage.OutputTokens != 5 {
		t.Errorf("usage tokens: got %+v", resp.Usage)
	}
	wantCost := float64(10)/1_000_000*0.30 + float64(5)/1_000_000*2.50
	if resp.Usage.CostUSD != wantCost {
		t.Errorf("CostUSD: want %v got %v", wantCost, resp.Usage.CostUSD)
	}

	// Verify the outgoing request shape: system is a top-level
	// systemInstruction Content, not a message in "contents".
	if got.SystemInstruction == nil || len(got.SystemInstruction.Parts) != 1 || got.SystemInstruction.Parts[0].Text != "You are helpful." {
		t.Errorf("systemInstruction: got %+v", got.SystemInstruction)
	}
	if len(got.Contents) != 1 || got.Contents[0].Role != "user" {
		t.Fatalf("contents: %+v", got.Contents)
	}
	if len(got.Contents[0].Parts) != 1 || got.Contents[0].Parts[0].Text != "hi" {
		t.Errorf("user content part: %+v", got.Contents[0].Parts)
	}
}

// TestChat_ToolUseResponse verifies that a response whose parts mix text and
// a functionCall round-trips into ChatResponse.ToolCalls correctly
// (id/name/args), and that tools declared on the request are translated to
// Gemini's {tools: [{functionDeclarations: [...]}]} shape with "parameters"
// (not "input_schema" or a "function" wrapper).
func TestChat_ToolUseResponse(t *testing.T) {
	var got wireRequest
	srv := fakeGoogleServer(t, http.StatusOK, `{
		"candidates": [{
			"content": {
				"role": "model",
				"parts": [
					{"text": "Let me check that."},
					{"functionCall": {"id": "call_01", "name": "read_file", "args": {"path": "a.txt"}}}
				]
			},
			"finishReason": "STOP"
		}],
		"usageMetadata": {"promptTokenCount": 20, "candidatesTokenCount": 8, "totalTokenCount": 28}
	}`, &got)
	defer srv.Close()

	p := newTestProvider(srv.URL)
	resp, err := p.Chat(context.Background(), provider.ChatRequest{
		Model:    "gemini-2.5-pro",
		Messages: []provider.Message{{Role: "user", Text: "read a.txt"}},
		Tools: []provider.ToolSpec{
			{
				Name:        "read_file",
				Description: "Read a file.",
				ParametersJSONSchema: map[string]any{
					"type":       "object",
					"properties": map[string]any{"path": map[string]any{"type": "string"}},
					"required":   []string{"path"},
				},
			},
		},
	})
	if err != nil {
		t.Fatalf("Chat: %v", err)
	}
	if resp.Text != "Let me check that." {
		t.Errorf("Text: got %q", resp.Text)
	}
	if resp.StopReason != "STOP" {
		t.Errorf("StopReason: want STOP got %q", resp.StopReason)
	}
	if len(resp.ToolCalls) != 1 {
		t.Fatalf("expected 1 tool call, got %d: %+v", len(resp.ToolCalls), resp.ToolCalls)
	}
	tc := resp.ToolCalls[0]
	if tc.ID != "call_01" || tc.Name != "read_file" {
		t.Errorf("tool call id/name: got %+v", tc)
	}
	var args struct {
		Path string `json:"path"`
	}
	if err := json.Unmarshal([]byte(tc.ArgsJSON), &args); err != nil {
		t.Fatalf("tool call args not valid JSON: %v (%q)", err, tc.ArgsJSON)
	}
	if args.Path != "a.txt" {
		t.Errorf("tool call args.path: got %q", args.Path)
	}

	// Verify tools were declared as tools:[{functionDeclarations:[...]}].
	if len(got.Tools) != 1 || len(got.Tools[0].FunctionDeclarations) != 1 {
		t.Fatalf("expected 1 tool wrapper with 1 declaration, got %+v", got.Tools)
	}
	decl := got.Tools[0].FunctionDeclarations[0]
	if decl.Name != "read_file" || decl.Description != "Read a file." {
		t.Errorf("function declaration: %+v", decl)
	}
	if decl.Parameters == nil || decl.Parameters["type"] != "object" {
		t.Errorf("parameters not forwarded: %+v", decl.Parameters)
	}
}

// TestChat_MultiTurnWithPriorToolResult exercises the functionResponse-as-
// a-user-content translation: a provider.Message carrying ToolResults (as
// agentloop.Loop emits with Role "tool" after dispatching a tool call) must
// serialize as a **user**-role Content with functionResponse parts — Gemini
// has no "tool"/"function" role. Also verifies a prior "assistant" turn
// with a tool call round-trips into a single "model"-role Content with a
// functionCall part carrying the same id used in the subsequent
// functionResponse.
func TestChat_MultiTurnWithPriorToolResult(t *testing.T) {
	var got wireRequest
	srv := fakeGoogleServer(t, http.StatusOK, `{
		"candidates": [{
			"content": {"role": "model", "parts": [{"text": "The file says hello."}]},
			"finishReason": "STOP"
		}],
		"usageMetadata": {"promptTokenCount": 30, "candidatesTokenCount": 6, "totalTokenCount": 36}
	}`, &got)
	defer srv.Close()

	p := newTestProvider(srv.URL)
	_, err := p.Chat(context.Background(), provider.ChatRequest{
		Model: "gemini-2.5-flash",
		Messages: []provider.Message{
			{Role: "user", Text: "read a.txt"},
			{
				Role: "assistant",
				Text: "",
				ToolCalls: []provider.ToolCall{
					{ID: "call_01", Name: "read_file", ArgsJSON: `{"path":"a.txt"}`},
				},
			},
			{
				Role: "tool",
				ToolResults: []provider.ToolResult{
					{ToolCallID: "call_01", Name: "read_file", Content: "hello", IsError: false},
				},
			},
		},
	})
	if err != nil {
		t.Fatalf("Chat: %v", err)
	}

	if len(got.Contents) != 3 {
		t.Fatalf("expected 3 wire contents, got %d: %+v", len(got.Contents), got.Contents)
	}

	// Turn 2: model functionCall, no text part since Text was empty.
	asst := got.Contents[1]
	if asst.Role != "model" {
		t.Errorf("turn 2 role: want model got %q", asst.Role)
	}
	if len(asst.Parts) != 1 || asst.Parts[0].FunctionCall == nil {
		t.Fatalf("turn 2 parts: %+v", asst.Parts)
	}
	if asst.Parts[0].FunctionCall.ID != "call_01" || asst.Parts[0].FunctionCall.Name != "read_file" {
		t.Errorf("turn 2 functionCall part: %+v", asst.Parts[0].FunctionCall)
	}

	// Turn 3: the tool result must be a USER-role Content with a
	// functionResponse part (Gemini has no "tool"/"function" role).
	toolTurn := got.Contents[2]
	if toolTurn.Role != "user" {
		t.Errorf("tool result turn role: want user (Gemini has no tool role) got %q", toolTurn.Role)
	}
	if len(toolTurn.Parts) != 1 || toolTurn.Parts[0].FunctionResponse == nil {
		t.Fatalf("tool result parts: %+v", toolTurn.Parts)
	}
	fr := toolTurn.Parts[0].FunctionResponse
	if fr.ID != "call_01" {
		t.Errorf("functionResponse id: want call_01 got %q", fr.ID)
	}
	if fr.Name != "read_file" {
		t.Errorf("functionResponse name: want read_file got %q", fr.Name)
	}
	if fr.Response["content"] != "hello" {
		t.Errorf("functionResponse.response.content: want hello got %+v", fr.Response)
	}
	if _, hasErr := fr.Response["error"]; hasErr {
		t.Errorf("response should not carry an error key when IsError is false: %+v", fr.Response)
	}
}

// TestChat_ToolResultIsError verifies IsError round-trips onto the
// functionResponse's response object as an "error" key (this adapter's
// convention — Gemini imposes no schema on the response object).
func TestChat_ToolResultIsError(t *testing.T) {
	var got wireRequest
	srv := fakeGoogleServer(t, http.StatusOK, `{
		"candidates": [{"content": {"role": "model", "parts": [{"text": "ok"}]}, "finishReason": "STOP"}],
		"usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2}
	}`, &got)
	defer srv.Close()

	p := newTestProvider(srv.URL)
	_, err := p.Chat(context.Background(), provider.ChatRequest{
		Model: "gemini-2.5-flash",
		Messages: []provider.Message{
			{Role: "tool", ToolResults: []provider.ToolResult{
				{ToolCallID: "call_x", Name: "run_bash", Content: "boom", IsError: true},
			}},
		},
	})
	if err != nil {
		t.Fatalf("Chat: %v", err)
	}
	if len(got.Contents) != 1 || len(got.Contents[0].Parts) != 1 {
		t.Fatalf("unexpected contents: %+v", got.Contents)
	}
	fr := got.Contents[0].Parts[0].FunctionResponse
	if fr == nil {
		t.Fatalf("expected a functionResponse part")
	}
	if fr.Response["error"] != "boom" {
		t.Errorf("expected response.error = boom, got %+v", fr.Response)
	}
}

// TestChat_RateLimitError429 verifies a 429 RESOURCE_EXHAUSTED response
// produces a Go error whose text retry.IsRateLimitError recognizes.
func TestChat_RateLimitError429(t *testing.T) {
	srv := fakeGoogleServer(t, http.StatusTooManyRequests, `{
		"error": {"code": 429, "message": "You exceeded your current quota.", "status": "RESOURCE_EXHAUSTED"}
	}`, nil)
	defer srv.Close()

	p := newTestProvider(srv.URL)
	_, err := p.Chat(context.Background(), provider.ChatRequest{
		Model:    "gemini-2.5-flash",
		Messages: []provider.Message{{Role: "user", Text: "hi"}},
	})
	if err == nil {
		t.Fatal("expected an error for HTTP 429")
	}
	if !strings.Contains(err.Error(), "429") || !strings.Contains(err.Error(), "RESOURCE_EXHAUSTED") {
		t.Errorf("error should mention 429 and RESOURCE_EXHAUSTED, got: %v", err)
	}
	if !retry.IsRateLimitError(err) {
		t.Errorf("retry.IsRateLimitError should recognize: %v", err)
	}
}

// TestChat_InvalidArgumentError400 is a non-retryable error case: confirms
// the error text carries the INVALID_ARGUMENT status and is NOT
// misclassified as a rate-limit error.
func TestChat_InvalidArgumentError400(t *testing.T) {
	srv := fakeGoogleServer(t, http.StatusBadRequest, `{
		"error": {"code": 400, "message": "Request contains an invalid argument.", "status": "INVALID_ARGUMENT"}
	}`, nil)
	defer srv.Close()

	p := newTestProvider(srv.URL)
	_, err := p.Chat(context.Background(), provider.ChatRequest{
		Model:    "gemini-2.5-flash",
		Messages: []provider.Message{{Role: "user", Text: "hi"}},
	})
	if err == nil {
		t.Fatal("expected an error for HTTP 400")
	}
	if !strings.Contains(err.Error(), "INVALID_ARGUMENT") {
		t.Errorf("error should mention INVALID_ARGUMENT, got: %v", err)
	}
	if retry.IsRateLimitError(err) {
		t.Errorf("400 INVALID_ARGUMENT should NOT be classified as a rate-limit error: %v", err)
	}
}

// TestChat_NoAPIKey_Errors confirms Chat fails fast without ever making an
// HTTP request when no API key is configured.
func TestChat_NoAPIKey_Errors(t *testing.T) {
	p := New("", "http://127.0.0.1:1", time.Second)
	_, err := p.Chat(context.Background(), provider.ChatRequest{Model: "gemini-2.5-flash"})
	if err == nil || !strings.Contains(err.Error(), "no API key") {
		t.Errorf("expected 'no API key' error, got %v", err)
	}
}

// TestChat_NoModel_Errors confirms Chat fails fast when no model is set,
// since the model name is part of the URL path (unlike Anthropic, where a
// missing model just gets sent through and rejected server-side).
func TestChat_NoModel_Errors(t *testing.T) {
	p := New("test-key", "http://127.0.0.1:1", time.Second)
	_, err := p.Chat(context.Background(), provider.ChatRequest{})
	if err == nil || !strings.Contains(err.Error(), "no model") {
		t.Errorf("expected 'no model' error, got %v", err)
	}
}

// TestNew_DefaultsEndpointAndTimeout verifies New falls back to
// DefaultEndpoint/DefaultTimeout for zero-value arguments.
func TestNew_DefaultsEndpointAndTimeout(t *testing.T) {
	p := New("key", "", 0)
	if p.Endpoint != DefaultEndpoint {
		t.Errorf("Endpoint: want %q got %q", DefaultEndpoint, p.Endpoint)
	}
	if p.HTTPClient.Timeout != DefaultTimeout {
		t.Errorf("Timeout: want %v got %v", DefaultTimeout, p.HTTPClient.Timeout)
	}
}

// TestName returns the provider's identifier, used as the agent-type/budget
// key ("google") once wired into a NativeRunner.
func TestName(t *testing.T) {
	p := New("key", "", 0)
	if p.Name() != "google" {
		t.Errorf("Name: got %q", p.Name())
	}
}

// TestPricingFor_KnownAndUnknownModels sanity-checks the pricing table
// lookup: exact/prefix match for known models, and a non-erroring fallback
// for an unrecognized model string.
func TestPricingFor_KnownAndUnknownModels(t *testing.T) {
	in, out := pricingFor("gemini-2.5-pro")
	if in != 1.25 || out != 10.00 {
		t.Errorf("2.5 pro pricing: got %v/%v", in, out)
	}
	in, out = pricingFor("gemini-2.5-flash-lite-preview-09-2025")
	if in != 0.10 || out != 0.40 {
		t.Errorf("2.5 flash-lite dated snapshot pricing: got %v/%v", in, out)
	}
	in, out = pricingFor("gemini-2.5-flash-preview")
	if in != 0.30 || out != 2.50 {
		t.Errorf("2.5 flash (not lite) pricing: got %v/%v", in, out)
	}
	in, out = pricingFor("some-future-unknown-model")
	if in != defaultPricing.InputPerMTok || out != defaultPricing.OutputPerMTok {
		t.Errorf("unknown model should fall back to default pricing, got %v/%v", in, out)
	}
}