Billy D.
35912d5844
- messages/bench_test.go: serialization benchmarks (msgpack map vs struct vs protobuf) - clients/clients_test.go: HTTP client tests with pooling verification (20 tests) - natsutil/natsutil_test.go: encode/decode roundtrip + binary data tests - handler/handler_test.go: handler dispatch tests + benchmark - config/config.go: live reload via fsnotify + RWMutex getter methods - clients/clients.go: SharedTransport + sync.Pool buffer pooling - messages/messages.go: typed structs with msgpack+json tags - messages/proto/: protobuf schema + generated code Benchmark baseline (ChatRequest roundtrip): MsgpackMap: 2949 ns/op, 36 allocs MsgpackStruct: 2030 ns/op, 13 allocs (31% faster, 64% fewer allocs) Protobuf: 793 ns/op, 8 allocs (73% faster, 78% fewer allocs)
507 lines
16 KiB
Go
507 lines
16 KiB
Go
package clients
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
// ────────────────────────────────────────────────────────────────────────────
|
|
// Shared infrastructure tests
|
|
// ────────────────────────────────────────────────────────────────────────────
|
|
|
|
func TestSharedTransport(t *testing.T) {
|
|
// All clients created via newHTTPClient should share the same transport.
|
|
c1 := newHTTPClient("http://a:8000", 10*time.Second)
|
|
c2 := newHTTPClient("http://b:9000", 30*time.Second)
|
|
|
|
if c1.client.Transport != c2.client.Transport {
|
|
t.Error("clients should share the same http.Transport")
|
|
}
|
|
if c1.client.Transport != SharedTransport {
|
|
t.Error("transport should be the package-level SharedTransport")
|
|
}
|
|
}
|
|
|
|
func TestBufferPoolGetPut(t *testing.T) {
|
|
buf := getBuf()
|
|
if buf == nil {
|
|
t.Fatal("getBuf returned nil")
|
|
}
|
|
if buf.Len() != 0 {
|
|
t.Error("getBuf should return a reset buffer")
|
|
}
|
|
buf.WriteString("hello")
|
|
putBuf(buf)
|
|
|
|
// On re-get, buffer should be reset.
|
|
buf2 := getBuf()
|
|
if buf2.Len() != 0 {
|
|
t.Error("re-acquired buffer should be reset")
|
|
}
|
|
putBuf(buf2)
|
|
}
|
|
|
|
func TestBufferPoolOversizedDiscarded(t *testing.T) {
|
|
buf := getBuf()
|
|
// Grow beyond 1 MB threshold.
|
|
buf.Write(make([]byte, 2<<20))
|
|
putBuf(buf) // should silently discard
|
|
|
|
// Pool should still work — we get a fresh one.
|
|
buf2 := getBuf()
|
|
if buf2.Len() != 0 {
|
|
t.Error("should get a fresh buffer")
|
|
}
|
|
putBuf(buf2)
|
|
}
|
|
|
|
func TestBufferPoolConcurrency(t *testing.T) {
|
|
var wg sync.WaitGroup
|
|
for i := range 100 {
|
|
wg.Add(1)
|
|
go func(n int) {
|
|
defer wg.Done()
|
|
buf := getBuf()
|
|
buf.WriteString(strings.Repeat("x", n))
|
|
putBuf(buf)
|
|
}(i)
|
|
}
|
|
wg.Wait()
|
|
}
|
|
|
|
// ────────────────────────────────────────────────────────────────────────────
|
|
// Embeddings client
|
|
// ────────────────────────────────────────────────────────────────────────────
|
|
|
|
func TestEmbeddingsClient_Embed(t *testing.T) {
|
|
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path != "/embeddings" {
|
|
t.Errorf("path = %q, want /embeddings", r.URL.Path)
|
|
}
|
|
if r.Method != http.MethodPost {
|
|
t.Errorf("method = %s, want POST", r.Method)
|
|
}
|
|
var req map[string]any
|
|
json.NewDecoder(r.Body).Decode(&req)
|
|
input, _ := req["input"].([]any)
|
|
if len(input) != 2 {
|
|
t.Errorf("input len = %d, want 2", len(input))
|
|
}
|
|
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(map[string]any{
|
|
"data": []map[string]any{
|
|
{"embedding": []float64{0.1, 0.2, 0.3}},
|
|
{"embedding": []float64{0.4, 0.5, 0.6}},
|
|
},
|
|
})
|
|
}))
|
|
defer ts.Close()
|
|
|
|
c := NewEmbeddingsClient(ts.URL, 5*time.Second, "bge")
|
|
results, err := c.Embed(context.Background(), []string{"hello", "world"})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(results) != 2 {
|
|
t.Fatalf("len(results) = %d, want 2", len(results))
|
|
}
|
|
if results[0][0] != 0.1 {
|
|
t.Errorf("results[0][0] = %f, want 0.1", results[0][0])
|
|
}
|
|
}
|
|
|
|
func TestEmbeddingsClient_EmbedSingle(t *testing.T) {
|
|
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
json.NewEncoder(w).Encode(map[string]any{
|
|
"data": []map[string]any{
|
|
{"embedding": []float64{1.0, 2.0}},
|
|
},
|
|
})
|
|
}))
|
|
defer ts.Close()
|
|
|
|
c := NewEmbeddingsClient(ts.URL, 5*time.Second, "")
|
|
vec, err := c.EmbedSingle(context.Background(), "test")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(vec) != 2 || vec[0] != 1.0 {
|
|
t.Errorf("vec = %v", vec)
|
|
}
|
|
}
|
|
|
|
func TestEmbeddingsClient_EmbedEmpty(t *testing.T) {
|
|
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
json.NewEncoder(w).Encode(map[string]any{"data": []any{}})
|
|
}))
|
|
defer ts.Close()
|
|
|
|
c := NewEmbeddingsClient(ts.URL, 5*time.Second, "")
|
|
_, err := c.EmbedSingle(context.Background(), "test")
|
|
if err == nil {
|
|
t.Error("expected error for empty embedding")
|
|
}
|
|
}
|
|
|
|
func TestEmbeddingsClient_Health(t *testing.T) {
|
|
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path == "/health" {
|
|
w.WriteHeader(200)
|
|
return
|
|
}
|
|
w.WriteHeader(404)
|
|
}))
|
|
defer ts.Close()
|
|
|
|
c := NewEmbeddingsClient(ts.URL, 5*time.Second, "")
|
|
if !c.Health(context.Background()) {
|
|
t.Error("expected healthy")
|
|
}
|
|
}
|
|
|
|
// ────────────────────────────────────────────────────────────────────────────
|
|
// Reranker client
|
|
// ────────────────────────────────────────────────────────────────────────────
|
|
|
|
func TestRerankerClient_Rerank(t *testing.T) {
|
|
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
var req map[string]any
|
|
json.NewDecoder(r.Body).Decode(&req)
|
|
if req["query"] != "test query" {
|
|
t.Errorf("query = %v", req["query"])
|
|
}
|
|
json.NewEncoder(w).Encode(map[string]any{
|
|
"results": []map[string]any{
|
|
{"index": 1, "relevance_score": 0.95},
|
|
{"index": 0, "relevance_score": 0.80},
|
|
},
|
|
})
|
|
}))
|
|
defer ts.Close()
|
|
|
|
c := NewRerankerClient(ts.URL, 5*time.Second)
|
|
docs := []string{"Paris is great", "France is in Europe"}
|
|
results, err := c.Rerank(context.Background(), "test query", docs, 2)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(results) != 2 {
|
|
t.Fatalf("len = %d", len(results))
|
|
}
|
|
if results[0].Score != 0.95 {
|
|
t.Errorf("score = %f, want 0.95", results[0].Score)
|
|
}
|
|
if results[0].Document != "France is in Europe" {
|
|
t.Errorf("document = %q", results[0].Document)
|
|
}
|
|
}
|
|
|
|
func TestRerankerClient_RerankFallbackScore(t *testing.T) {
|
|
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
json.NewEncoder(w).Encode(map[string]any{
|
|
"results": []map[string]any{
|
|
{"index": 0, "score": 0.77, "relevance_score": 0}, // some APIs only set score
|
|
},
|
|
})
|
|
}))
|
|
defer ts.Close()
|
|
|
|
c := NewRerankerClient(ts.URL, 5*time.Second)
|
|
results, err := c.Rerank(context.Background(), "q", []string{"doc1"}, 0)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if results[0].Score != 0.77 {
|
|
t.Errorf("fallback score = %f, want 0.77", results[0].Score)
|
|
}
|
|
}
|
|
|
|
// ────────────────────────────────────────────────────────────────────────────
|
|
// LLM client
|
|
// ────────────────────────────────────────────────────────────────────────────
|
|
|
|
func TestLLMClient_Generate(t *testing.T) {
|
|
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path != "/v1/chat/completions" {
|
|
t.Errorf("path = %q", r.URL.Path)
|
|
}
|
|
var req map[string]any
|
|
json.NewDecoder(r.Body).Decode(&req)
|
|
msgs, _ := req["messages"].([]any)
|
|
if len(msgs) == 0 {
|
|
t.Error("no messages in request")
|
|
}
|
|
|
|
json.NewEncoder(w).Encode(map[string]any{
|
|
"choices": []map[string]any{
|
|
{"message": map[string]any{"content": "Paris is the capital of France."}},
|
|
},
|
|
})
|
|
}))
|
|
defer ts.Close()
|
|
|
|
c := NewLLMClient(ts.URL, 5*time.Second)
|
|
result, err := c.Generate(context.Background(), "capital of France?", "", "")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if result != "Paris is the capital of France." {
|
|
t.Errorf("result = %q", result)
|
|
}
|
|
}
|
|
|
|
func TestLLMClient_GenerateWithContext(t *testing.T) {
|
|
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
var req map[string]any
|
|
json.NewDecoder(r.Body).Decode(&req)
|
|
msgs, _ := req["messages"].([]any)
|
|
// Should have system + user message
|
|
if len(msgs) != 2 {
|
|
t.Errorf("expected 2 messages, got %d", len(msgs))
|
|
}
|
|
json.NewEncoder(w).Encode(map[string]any{
|
|
"choices": []map[string]any{
|
|
{"message": map[string]any{"content": "answer with context"}},
|
|
},
|
|
})
|
|
}))
|
|
defer ts.Close()
|
|
|
|
c := NewLLMClient(ts.URL, 5*time.Second)
|
|
result, err := c.Generate(context.Background(), "question", "some context", "")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if result != "answer with context" {
|
|
t.Errorf("result = %q", result)
|
|
}
|
|
}
|
|
|
|
func TestLLMClient_GenerateNoChoices(t *testing.T) {
|
|
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
json.NewEncoder(w).Encode(map[string]any{"choices": []any{}})
|
|
}))
|
|
defer ts.Close()
|
|
|
|
c := NewLLMClient(ts.URL, 5*time.Second)
|
|
_, err := c.Generate(context.Background(), "q", "", "")
|
|
if err == nil {
|
|
t.Error("expected error for empty choices")
|
|
}
|
|
}
|
|
|
|
// ────────────────────────────────────────────────────────────────────────────
|
|
// TTS client
|
|
// ────────────────────────────────────────────────────────────────────────────
|
|
|
|
func TestTTSClient_Synthesize(t *testing.T) {
|
|
expected := []byte{0xDE, 0xAD, 0xBE, 0xEF}
|
|
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path != "/api/tts" {
|
|
t.Errorf("path = %q", r.URL.Path)
|
|
}
|
|
if r.URL.Query().Get("text") != "hello world" {
|
|
t.Errorf("text = %q", r.URL.Query().Get("text"))
|
|
}
|
|
w.Write(expected)
|
|
}))
|
|
defer ts.Close()
|
|
|
|
c := NewTTSClient(ts.URL, 5*time.Second, "en")
|
|
audio, err := c.Synthesize(context.Background(), "hello world", "", "")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if !bytes.Equal(audio, expected) {
|
|
t.Errorf("audio = %x, want %x", audio, expected)
|
|
}
|
|
}
|
|
|
|
func TestTTSClient_SynthesizeWithSpeaker(t *testing.T) {
|
|
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Query().Get("speaker_id") != "alice" {
|
|
t.Errorf("speaker_id = %q", r.URL.Query().Get("speaker_id"))
|
|
}
|
|
w.Write([]byte{0x01})
|
|
}))
|
|
defer ts.Close()
|
|
|
|
c := NewTTSClient(ts.URL, 5*time.Second, "en")
|
|
_, err := c.Synthesize(context.Background(), "hi", "en", "alice")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
|
|
// ────────────────────────────────────────────────────────────────────────────
|
|
// STT client
|
|
// ────────────────────────────────────────────────────────────────────────────
|
|
|
|
func TestSTTClient_Transcribe(t *testing.T) {
|
|
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path != "/v1/audio/transcriptions" {
|
|
t.Errorf("path = %q", r.URL.Path)
|
|
}
|
|
ct := r.Header.Get("Content-Type")
|
|
if !strings.Contains(ct, "multipart/form-data") {
|
|
t.Errorf("content-type = %q", ct)
|
|
}
|
|
// Verify the audio file is present.
|
|
file, _, err := r.FormFile("file")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
data, _ := io.ReadAll(file)
|
|
if len(data) != 100 {
|
|
t.Errorf("file size = %d, want 100", len(data))
|
|
}
|
|
|
|
json.NewEncoder(w).Encode(map[string]string{"text": "hello world"})
|
|
}))
|
|
defer ts.Close()
|
|
|
|
c := NewSTTClient(ts.URL, 5*time.Second)
|
|
result, err := c.Transcribe(context.Background(), make([]byte, 100), "en")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if result.Text != "hello world" {
|
|
t.Errorf("text = %q", result.Text)
|
|
}
|
|
}
|
|
|
|
func TestSTTClient_TranscribeTranslate(t *testing.T) {
|
|
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path != "/v1/audio/translations" {
|
|
t.Errorf("path = %q, want /v1/audio/translations", r.URL.Path)
|
|
}
|
|
json.NewEncoder(w).Encode(map[string]string{"text": "translated"})
|
|
}))
|
|
defer ts.Close()
|
|
|
|
c := NewSTTClient(ts.URL, 5*time.Second)
|
|
c.Task = "translate"
|
|
result, err := c.Transcribe(context.Background(), []byte{0x01}, "")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if result.Text != "translated" {
|
|
t.Errorf("text = %q", result.Text)
|
|
}
|
|
}
|
|
|
|
// ────────────────────────────────────────────────────────────────────────────
|
|
// HTTP error handling
|
|
// ────────────────────────────────────────────────────────────────────────────
|
|
|
|
func TestHTTPError4xx(t *testing.T) {
|
|
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.WriteHeader(422)
|
|
w.Write([]byte(`{"error": "bad input"}`))
|
|
}))
|
|
defer ts.Close()
|
|
|
|
c := NewEmbeddingsClient(ts.URL, 5*time.Second, "")
|
|
_, err := c.Embed(context.Background(), []string{"test"})
|
|
if err == nil {
|
|
t.Fatal("expected error for 422")
|
|
}
|
|
if !strings.Contains(err.Error(), "422") {
|
|
t.Errorf("error should contain status code: %v", err)
|
|
}
|
|
}
|
|
|
|
func TestHTTPError5xx(t *testing.T) {
|
|
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.WriteHeader(500)
|
|
w.Write([]byte("internal server error"))
|
|
}))
|
|
defer ts.Close()
|
|
|
|
c := NewLLMClient(ts.URL, 5*time.Second)
|
|
_, err := c.Generate(context.Background(), "q", "", "")
|
|
if err == nil {
|
|
t.Fatal("expected error for 500")
|
|
}
|
|
}
|
|
|
|
// ────────────────────────────────────────────────────────────────────────────
|
|
// buildMessages helper
|
|
// ────────────────────────────────────────────────────────────────────────────
|
|
|
|
func TestBuildMessages(t *testing.T) {
|
|
// No context, no system prompt → just user message
|
|
msgs := buildMessages("hello", "", "")
|
|
if len(msgs) != 1 || msgs[0].Role != "user" {
|
|
t.Errorf("expected 1 user msg, got %+v", msgs)
|
|
}
|
|
|
|
// With system prompt
|
|
msgs = buildMessages("hello", "", "You are helpful")
|
|
if len(msgs) != 2 || msgs[0].Role != "system" || msgs[0].Content != "You are helpful" {
|
|
t.Errorf("expected system+user, got %+v", msgs)
|
|
}
|
|
|
|
// With context, no system prompt → auto system prompt
|
|
msgs = buildMessages("question", "some context", "")
|
|
if len(msgs) != 2 || msgs[0].Role != "system" {
|
|
t.Errorf("expected auto system+user, got %+v", msgs)
|
|
}
|
|
if !strings.Contains(msgs[1].Content, "Context:") {
|
|
t.Error("user message should contain context")
|
|
}
|
|
}
|
|
|
|
// ────────────────────────────────────────────────────────────────────────────
|
|
// Benchmarks: pooled buffer vs direct allocation
|
|
// ────────────────────────────────────────────────────────────────────────────
|
|
|
|
func BenchmarkPostJSON(b *testing.B) {
|
|
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
io.Copy(io.Discard, r.Body)
|
|
w.Write([]byte(`{"ok":true}`))
|
|
}))
|
|
defer ts.Close()
|
|
|
|
c := newHTTPClient(ts.URL, 10*time.Second)
|
|
ctx := context.Background()
|
|
payload := map[string]any{
|
|
"text": strings.Repeat("x", 1024),
|
|
"count": 42,
|
|
"enabled": true,
|
|
}
|
|
|
|
b.ResetTimer()
|
|
for b.Loop() {
|
|
c.postJSON(ctx, "/test", payload)
|
|
}
|
|
}
|
|
|
|
func BenchmarkBufferPool(b *testing.B) {
|
|
b.ResetTimer()
|
|
for b.Loop() {
|
|
buf := getBuf()
|
|
buf.WriteString(strings.Repeat("x", 4096))
|
|
putBuf(buf)
|
|
}
|
|
}
|
|
|
|
func BenchmarkBufferPoolParallel(b *testing.B) {
|
|
b.RunParallel(func(pb *testing.PB) {
|
|
for pb.Next() {
|
|
buf := getBuf()
|
|
buf.WriteString(strings.Repeat("x", 4096))
|
|
putBuf(buf)
|
|
}
|
|
})
|
|
}
|