GopherGate/internal/providers/gemini_test.go

package providers

import (
	"encoding/json"
	"reflect"
	"strings"
	"testing"

	"gophergate/internal/models"
)

func TestCleanParametersSchema(t *testing.T) {
	tests := []struct {
		name     string
		input    string
		expected string
	}{
		{
			name: "flat additionalProperties and $schema",
			input: `{
				"$schema": "http://json-schema.org/draft-07/schema#",
				"type": "object",
				"properties": {
					"query": {
						"type": "string"
					}
				},
				"required": ["query"],
				"additionalProperties": false
			}`,
			expected: `{
				"type": "object",
				"properties": {
					"query": {
						"type": "string"
					}
				},
				"required": ["query"]
			}`,
		},
		{
			name: "nested additionalProperties",
			input: `{
				"type": "object",
				"properties": {
					"user": {
						"type": "object",
						"properties": {
							"id": {
								"type": "integer"
							}
						},
						"additionalProperties": false
					}
				},
				"additionalProperties": true
			}`,
			expected: `{
				"type": "object",
				"properties": {
					"user": {
						"type": "object",
						"properties": {
							"id": {
								"type": "integer"
							}
						}
					}
				}
			}`,
		},
		{
			name: "nested additionalProperties in array items",
			input: `{
				"type": "object",
				"properties": {
					"items": {
						"type": "array",
						"items": {
							"type": "object",
							"properties": {
								"name": {
									"type": "string"
								}
							},
							"additionalProperties": false
						}
					}
				}
			}`,
			expected: `{
				"type": "object",
				"properties": {
					"items": {
						"type": "array",
						"items": {
							"type": "object",
							"properties": {
								"name": {
									"type": "string"
								}
							}
						}
					}
				}
			}`,
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			inputRaw := json.RawMessage(tc.input)
			cleanedRaw := cleanParametersSchema(inputRaw)

			var cleanedMap, expectedMap map[string]interface{}
			if err := json.Unmarshal(cleanedRaw, &cleanedMap); err != nil {
				t.Fatalf("failed to unmarshal cleaned: %v", err)
			}
			if err := json.Unmarshal([]byte(tc.expected), &expectedMap); err != nil {
				t.Fatalf("failed to unmarshal expected: %v", err)
			}

			if !reflect.DeepEqual(cleanedMap, expectedMap) {
				t.Errorf("expected %v, got %v", expectedMap, cleanedMap)
			}
		})
	}
}

func TestEmitGeminiChunk_ToolCalls(t *testing.T) {
	ch := make(chan *models.ChatCompletionStreamResponse, 1)
	defer close(ch)

	chunk := &geminiStreamChunk{}
	chunk.Candidates = []struct {
		Content struct {
			Parts []struct {
				Text         string `json:"text,omitempty"`
				Thought      string `json:"thought,omitempty"`
				FunctionCall *struct {
					Name string          `json:"name"`
					Args json.RawMessage `json:"args"`
				} `json:"functionCall,omitempty"`
			} `json:"parts"`
		} `json:"content"`
		FinishReason string `json:"finishReason"`
	}{
		{
			Content: struct {
				Parts []struct {
					Text         string `json:"text,omitempty"`
					Thought      string `json:"thought,omitempty"`
					FunctionCall *struct {
						Name string          `json:"name"`
						Args json.RawMessage `json:"args"`
					} `json:"functionCall,omitempty"`
				} `json:"parts"`
			}{
				Parts: []struct {
					Text         string `json:"text,omitempty"`
					Thought      string `json:"thought,omitempty"`
					FunctionCall *struct {
						Name string          `json:"name"`
						Args json.RawMessage `json:"args"`
					} `json:"functionCall,omitempty"`
				}{
					{
						FunctionCall: &struct {
							Name string          `json:"name"`
							Args json.RawMessage `json:"args"`
						}{
							Name: "google_search",
							Args: json.RawMessage(`{"query": "test"}`),
						},
					},
				},
			},
			FinishReason: "",
		},
	}

	emitted := emitGeminiChunk(ch, chunk, "gemini-3-flash-preview")
	if !emitted {
		t.Fatalf("expected emitGeminiChunk to return true")
	}

	select {
	case resp := <-ch:
		if len(resp.Choices) != 1 {
			t.Fatalf("expected 1 choice, got %d", len(resp.Choices))
		}
		choice := resp.Choices[0]
		if choice.FinishReason == nil || *choice.FinishReason != "tool_calls" {
			t.Errorf("expected finish_reason 'tool_calls', got %v", choice.FinishReason)
		}
		if len(choice.Delta.ToolCalls) != 1 {
			t.Fatalf("expected 1 tool call in delta, got %d", len(choice.Delta.ToolCalls))
		}
		tc := choice.Delta.ToolCalls[0]
		if tc.ID == nil || *tc.ID != "call_google_search" {
			t.Errorf("expected ID 'call_google_search', got %v", tc.ID)
		}
		if tc.Function == nil || tc.Function.Name == nil || *tc.Function.Name != "google_search" {
			t.Errorf("expected function name 'google_search', got %v", tc.Function)
		}
		if tc.Function == nil || tc.Function.Arguments == nil || *tc.Function.Arguments != `{"query": "test"}` {
			t.Errorf("expected arguments '{\"query\": \"test\"}', got %v", tc.Function)
		}
	default:
		t.Fatalf("expected response on channel")
	}
}

func TestStreamGeminiJSONArrayStream(t *testing.T) {
	input := `[
		{
			"candidates": [
				{
					"content": {
						"parts": [
							{"text": "Hello"}
						]
					}
				}
			]
		},
		{
			"candidates": [
				{
					"content": {
						"parts": [
							{"text": " world"}
						]
					}
				}
			],
			"usageMetadata": {
				"promptTokenCount": 5,
				"candidatesTokenCount": 2,
				"totalTokenCount": 7
			}
		}
	]`

	ch := make(chan *models.ChatCompletionStreamResponse, 5)
	defer close(ch)

	streamGeminiJSONArrayStream(strings.NewReader(input), ch, "gemini-2.5-flash")

	var texts []string
	for i := 0; i < 2; i++ {
		select {
		case resp := <-ch:
			if len(resp.Choices) > 0 && resp.Choices[0].Delta.Content != nil {
				texts = append(texts, *resp.Choices[0].Delta.Content)
			}
		default:
			t.Fatalf("expected chunk %d", i+1)
		}
	}

	joined := strings.Join(texts, "")
	if joined != "Hello world" {
		t.Errorf("expected 'Hello world', got %q", joined)
	}
}