Files
iasysadmin/internal/llm/client_test.go
T
2026-09-27 16:04:02 -06:00

139 lines
4.4 KiB
Go

package llm
import (
"context"
"encoding/json"
"io"
"net/http"
"net/http/httptest"
"strings"
"sync/atomic"
"testing"
"time"
)
func newTestClient(url, key string) *OpenAIClient {
c := NewOpenAIClient(Options{BaseURL: url + "/v1", APIKey: key, Model: "qwen3", UserAgent: "sysagent/test", Timeout: 5 * time.Second})
c.backoff = time.Millisecond
return c
}
func TestChat(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/v1/chat/completions" || r.Method != http.MethodPost {
t.Errorf("unexpected %s %s", r.Method, r.URL.Path)
}
if got := r.Header.Get("Authorization"); got != "Bearer sk-secret" {
t.Errorf("Authorization = %q", got)
}
if got := r.Header.Get("User-Agent"); got != "sysagent/test" {
t.Errorf("User-Agent = %q", got)
}
body, _ := io.ReadAll(r.Body)
var req map[string]any
if err := json.Unmarshal(body, &req); err != nil {
t.Fatal(err)
}
if req["model"] != "qwen3" || req["stream"] != false {
t.Errorf("request = %v", req)
}
io.WriteString(w, `{"choices":[{"message":{"role":"assistant","content":"<think>hmm</think>\nHola"},"finish_reason":"stop"}],
"usage":{"prompt_tokens":10,"completion_tokens":2,"total_tokens":12}}`)
}))
defer srv.Close()
resp, err := newTestClient(srv.URL, "sk-secret").Chat(context.Background(), ChatRequest{
Messages: []Message{{Role: RoleUser, Content: "hola"}},
})
if err != nil {
t.Fatal(err)
}
if resp.Message.Content != "Hola" {
t.Errorf("content = %q (think block must be stripped)", resp.Message.Content)
}
if resp.Reasoning != "hmm" {
t.Errorf("reasoning = %q", resp.Reasoning)
}
if resp.Usage.PromptTokens != 10 || resp.FinishReason != "stop" {
t.Errorf("resp = %+v", resp)
}
}
func TestChatNoAuthHeaderWhenKeyEmpty(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Header.Get("Authorization") != "" {
t.Error("Authorization header must be absent when API_KEY is empty")
}
io.WriteString(w, `{"choices":[{"message":{"content":"ok"}}]}`)
}))
defer srv.Close()
if _, err := newTestClient(srv.URL, "").Chat(context.Background(), ChatRequest{}); err != nil {
t.Fatal(err)
}
}
func TestChatRetriesTransient(t *testing.T) {
var calls atomic.Int32
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if calls.Add(1) < 3 {
http.Error(w, "busy", http.StatusServiceUnavailable)
return
}
io.WriteString(w, `{"choices":[{"message":{"content":"ok"}}]}`)
}))
defer srv.Close()
resp, err := newTestClient(srv.URL, "").Chat(context.Background(), ChatRequest{})
if err != nil || resp.Message.Content != "ok" || calls.Load() != 3 {
t.Fatalf("resp=%+v err=%v calls=%d", resp, err, calls.Load())
}
}
func TestChatDoesNotRetryClientError(t *testing.T) {
var calls atomic.Int32
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
calls.Add(1)
http.Error(w, `{"error":{"message":"bad model"}}`, http.StatusBadRequest)
}))
defer srv.Close()
_, err := newTestClient(srv.URL, "").Chat(context.Background(), ChatRequest{})
if err == nil || !strings.Contains(err.Error(), "HTTP 400") || calls.Load() != 1 {
t.Fatalf("err=%v calls=%d", err, calls.Load())
}
}
func TestDiscoverModel(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/v1/models":
io.WriteString(w, `{"data":[{"id":"other","meta":{"n_ctx_train":131072}}]}`)
case "/props":
io.WriteString(w, `{"default_generation_settings":{"n_ctx":16384}}`)
default:
http.NotFound(w, r)
}
}))
defer srv.Close()
info, err := newTestClient(srv.URL, "").DiscoverModel(context.Background())
if err != nil {
t.Fatal(err)
}
if info.ContextSize != 16384 || !strings.Contains(info.Source, "/props") {
t.Errorf("info = %+v (configured n_ctx must win over n_ctx_train)", info)
}
}
func TestDiscoverModelUnknown(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path == "/v1/models" {
io.WriteString(w, `{"data":[{"id":"qwen3"},{"id":"x"}]}`)
return
}
http.NotFound(w, r)
}))
defer srv.Close()
info, err := newTestClient(srv.URL, "").DiscoverModel(context.Background())
if err != nil || info.ContextSize != 0 {
t.Fatalf("expected UNKNOWN (0) without error, got %+v err=%v", info, err)
}
}