package llm import ( "context" "encoding/json" "net/http" "net/http/httptest" "strings" "testing" "git.totmin.ru/en2zmax/addrkit/internal/config" ) // chatServer returns a test server that echoes a fixed content and captures the // request body it received. func chatServer(t *testing.T, content string) (*httptest.Server, *chatRequest) { t.Helper() var got chatRequest ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _ = json.NewDecoder(r.Body).Decode(&got) w.Header().Set("Content-Type", "application/json") resp := map[string]any{"choices": []any{map[string]any{"message": map[string]any{"content": content}}}} _ = json.NewEncoder(w).Encode(resp) })) t.Cleanup(ts.Close) return ts, &got } func TestRewriteUsesConfiguredPrompt(t *testing.T) { ts, got := chatServer(t, `{"variants":["А"]}`) r := New(config.LLMConfig{URL: ts.URL, Model: "m", SystemPrompt: "CUSTOM PROMPT"}) out, err := r.Rewrite(context.Background(), "вход") if err != nil { t.Fatalf("Rewrite: %v", err) } if len(got.Messages) == 0 || got.Messages[0].Content != "CUSTOM PROMPT" { t.Errorf("system prompt not applied: %+v", got.Messages) } if len(out) != 1 || out[0] != "А" { t.Errorf("variants = %v", out) } } func TestRewriteFallsBackToEmbeddedPrompt(t *testing.T) { ts, got := chatServer(t, `{"variants":[]}`) r := New(config.LLMConfig{URL: ts.URL, Model: "m"}) if _, err := r.Rewrite(context.Background(), "x"); err != nil { t.Fatalf("Rewrite: %v", err) } if len(got.Messages) == 0 || !strings.Contains(got.Messages[0].Content, "нормализатор") { t.Errorf("embedded default prompt not used: %q", got.Messages[0].Content) } } func TestRewriteDisablesThinking(t *testing.T) { ts, got := chatServer(t, `{"variants":[]}`) r := New(config.LLMConfig{URL: ts.URL, Model: "m", DisableThinking: true}) if _, err := r.Rewrite(context.Background(), "x"); err != nil { t.Fatalf("Rewrite: %v", err) } if v, ok := got.ChatTemplateKwargs["enable_thinking"]; !ok || v != false { t.Errorf("enable_thinking not disabled: %+v", got.ChatTemplateKwargs) } } func TestParseVariants(t *testing.T) { cases := []struct { name string content string want int }{ {"clean", `{"variants":["a","b"]}`, 2}, {"with thinking", "hmm\n{\"variants\":[\"a\"]}", 1}, {"unterminated thinking", "hmm {\"variants\":[\"a\"]}", 0}, {"prose around", "Here you go: {\"variants\":[\"a\",\"b\",\"c\"]} done", 3}, {"garbage", "no json here", 0}, } for _, tt := range cases { t.Run(tt.name, func(t *testing.T) { if got := parseVariants(tt.content); len(got) != tt.want { t.Errorf("parseVariants(%q) = %v, want %d items", tt.content, got, tt.want) } }) } } func TestParseVariantsCapsCountAndLength(t *testing.T) { long := strings.Repeat("x", 600) content := `{"variants":["` + long + `","a","b","c","d"]}` out := parseVariants(content) if len(out) != 3 { t.Fatalf("variants = %v, want 3", out) } for _, v := range out { if len([]rune(v)) > 512 { t.Fatalf("oversized variant accepted: %d runes", len([]rune(v))) } } }