ScuttleBot

scuttlebot / internal / llm / gemini_test.go

Source Blame History 99 lines

5ac549c…	lmata	1	package llm
5ac549c…	lmata	2
5ac549c…	lmata	3	import (
5ac549c…	lmata	4	"context"
5ac549c…	lmata	5	"encoding/json"
5ac549c…	lmata	6	"net/http"
5ac549c…	lmata	7	"net/http/httptest"
5ac549c…	lmata	8	"testing"
5ac549c…	lmata	9	)
5ac549c…	lmata	10
5ac549c…	lmata	11	func TestGeminiSummarize(t *testing.T) {
5ac549c…	lmata	12	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
5ac549c…	lmata	13	if r.Method != "POST" {
5ac549c…	lmata	14	t.Errorf("expected POST request, got %s", r.Method)
5ac549c…	lmata	15	}
5ac549c…	lmata	16	if r.URL.Path != "/v1beta/models/gemini-1.5-flash:generateContent" {
5ac549c…	lmata	17	t.Errorf("unexpected path: %s", r.URL.Path)
5ac549c…	lmata	18	}
5ac549c…	lmata	19	if r.URL.Query().Get("key") != "test-api-key" {
5ac549c…	lmata	20	t.Errorf("expected api key test-api-key, got %s", r.URL.Query().Get("key"))
5ac549c…	lmata	21	}
5ac549c…	lmata	22
5ac549c…	lmata	23	resp := map[string]any{
5ac549c…	lmata	24	"candidates": []map[string]any{
5ac549c…	lmata	25	{
5ac549c…	lmata	26	"content": map[string]any{
5ac549c…	lmata	27	"parts": []map[string]any{
5ac549c…	lmata	28	{"text": "gemini response"},
5ac549c…	lmata	29	},
5ac549c…	lmata	30	},
5ac549c…	lmata	31	},
5ac549c…	lmata	32	},
5ac549c…	lmata	33	}
5ac549c…	lmata	34	_ = json.NewEncoder(w).Encode(resp)
5ac549c…	lmata	35	}))
5ac549c…	lmata	36	defer srv.Close()
5ac549c…	lmata	37
5ac549c…	lmata	38	p := newGeminiProvider(BackendConfig{
5ac549c…	lmata	39	Backend: "gemini",
5ac549c…	lmata	40	APIKey: "test-api-key",
5ac549c…	lmata	41	BaseURL: srv.URL,
5ac549c…	lmata	42	}, srv.Client())
5ac549c…	lmata	43
5ac549c…	lmata	44	got, err := p.Summarize(context.Background(), "test prompt")
5ac549c…	lmata	45	if err != nil {
5ac549c…	lmata	46	t.Fatalf("Summarize failed: %v", err)
5ac549c…	lmata	47	}
5ac549c…	lmata	48	if got != "gemini response" {
5ac549c…	lmata	49	t.Errorf("got %q, want %q", got, "gemini response")
5ac549c…	lmata	50	}
5ac549c…	lmata	51	}
5ac549c…	lmata	52
5ac549c…	lmata	53	func TestGeminiDiscoverModels(t *testing.T) {
5ac549c…	lmata	54	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
5ac549c…	lmata	55	if r.Method != "GET" {
5ac549c…	lmata	56	t.Errorf("expected GET request, got %s", r.Method)
5ac549c…	lmata	57	}
5ac549c…	lmata	58	if r.URL.Path != "/v1beta/models" {
5ac549c…	lmata	59	t.Errorf("unexpected path: %s", r.URL.Path)
5ac549c…	lmata	60	}
5ac549c…	lmata	61
5ac549c…	lmata	62	resp := map[string]any{
5ac549c…	lmata	63	"models": []map[string]any{
5ac549c…	lmata	64	{
1066004…	lmata	65	"name": "models/gemini-1.5-flash",
1066004…	lmata	66	"displayName": "Gemini 1.5 Flash",
1066004…	lmata	67	"description": "Fast and versatile",
5ac549c…	lmata	68	"supportedGenerationMethods": []string{"generateContent"},
5ac549c…	lmata	69	},
5ac549c…	lmata	70	{
1066004…	lmata	71	"name": "models/other-model",
1066004…	lmata	72	"displayName": "Other",
1066004…	lmata	73	"description": "Other model",
5ac549c…	lmata	74	"supportedGenerationMethods": []string{"other"},
5ac549c…	lmata	75	},
5ac549c…	lmata	76	},
5ac549c…	lmata	77	}
5ac549c…	lmata	78	_ = json.NewEncoder(w).Encode(resp)
5ac549c…	lmata	79	}))
5ac549c…	lmata	80	defer srv.Close()
5ac549c…	lmata	81
5ac549c…	lmata	82	p := newGeminiProvider(BackendConfig{
5ac549c…	lmata	83	Backend: "gemini",
5ac549c…	lmata	84	APIKey: "test-api-key",
5ac549c…	lmata	85	BaseURL: srv.URL,
5ac549c…	lmata	86	}, srv.Client())
5ac549c…	lmata	87
5ac549c…	lmata	88	models, err := p.DiscoverModels(context.Background())
5ac549c…	lmata	89	if err != nil {
5ac549c…	lmata	90	t.Fatalf("DiscoverModels failed: %v", err)
5ac549c…	lmata	91	}
5ac549c…	lmata	92
5ac549c…	lmata	93	if len(models) != 1 {
5ac549c…	lmata	94	t.Errorf("got %d models, want 1", len(models))
5ac549c…	lmata	95	}
5ac549c…	lmata	96	if models[0].ID != "gemini-1.5-flash" {
5ac549c…	lmata	97	t.Errorf("got ID %q, want %q", models[0].ID, "gemini-1.5-flash")
5ac549c…	lmata	98	}
5ac549c…	lmata	99	}

ScuttleBot

Keyboard Shortcuts