test: 添加各模块单元测试（config/llm/service/handler）
This commit is contained in:
279
backend/internal/llm/client_test.go
Normal file
279
backend/internal/llm/client_test.go
Normal file
@@ -0,0 +1,279 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"ai-agent-scaffold-go/internal/model"
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
// ============================================================
|
||||
// 辅助函数
|
||||
// ============================================================
|
||||
|
||||
func newTestClient(url string) *OpenAIClient {
|
||||
return NewOpenAIClient(url, "test-key", "gpt-4", 5*time.Second)
|
||||
}
|
||||
|
||||
// ============================================================
|
||||
// Generate 测试
|
||||
// ============================================================
|
||||
|
||||
func TestOpenAIClient_Generate_ReturnsContent(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
assert.Equal(t, "POST", r.Method)
|
||||
assert.Equal(t, "Bearer test-key", r.Header.Get("Authorization"))
|
||||
|
||||
json.NewEncoder(w).Encode(map[string]any{
|
||||
"choices": []map[string]any{
|
||||
{"message": map[string]any{"content": "hello world"}},
|
||||
},
|
||||
})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
client := newTestClient(srv.URL)
|
||||
reply, err := client.Generate(context.Background(), []model.ChatMessage{
|
||||
{Role: model.ChatRoleUser, Content: "hi"},
|
||||
}, nil)
|
||||
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, "hello world", reply.Content)
|
||||
assert.Empty(t, reply.ToolCalls)
|
||||
}
|
||||
|
||||
func TestOpenAIClient_Generate_ToolCalls(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
json.NewEncoder(w).Encode(map[string]any{
|
||||
"choices": []map[string]any{
|
||||
{"message": map[string]any{
|
||||
"content": "",
|
||||
"tool_calls": []map[string]any{
|
||||
{"id": "call_1", "type": "function", "function": map[string]any{
|
||||
"name": "search",
|
||||
"arguments": `{"query":"test"}`,
|
||||
}},
|
||||
},
|
||||
}},
|
||||
},
|
||||
})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
client := newTestClient(srv.URL)
|
||||
reply, err := client.Generate(context.Background(), []model.ChatMessage{
|
||||
{Role: model.ChatRoleUser, Content: "search for test"},
|
||||
}, []ToolDef{{Name: "search", Description: "search tool"}})
|
||||
|
||||
assert.NoError(t, err)
|
||||
assert.Len(t, reply.ToolCalls, 1)
|
||||
assert.Equal(t, "call_1", reply.ToolCalls[0].ID)
|
||||
assert.Equal(t, "search", reply.ToolCalls[0].Name)
|
||||
assert.Equal(t, `{"query":"test"}`, reply.ToolCalls[0].Arguments)
|
||||
}
|
||||
|
||||
func TestOpenAIClient_Generate_Upstream500_ReturnsError(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
fmt.Fprint(w, "internal error")
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
client := newTestClient(srv.URL)
|
||||
_, err := client.Generate(context.Background(), []model.ChatMessage{
|
||||
{Role: model.ChatRoleUser, Content: "hi"},
|
||||
}, nil)
|
||||
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "500")
|
||||
}
|
||||
|
||||
func TestOpenAIClient_Generate_EmptyChoices_ReturnsError(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
json.NewEncoder(w).Encode(map[string]any{"choices": []any{}})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
client := newTestClient(srv.URL)
|
||||
_, err := client.Generate(context.Background(), []model.ChatMessage{
|
||||
{Role: model.ChatRoleUser, Content: "hi"},
|
||||
}, nil)
|
||||
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "no choices")
|
||||
}
|
||||
|
||||
// ============================================================
|
||||
// Stream 测试
|
||||
// ============================================================
|
||||
|
||||
func TestOpenAIClient_Stream_ReturnsChunks(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
assert.Equal(t, "POST", r.Method)
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"hel\"}}]}\n\n")
|
||||
w.(http.Flusher).Flush()
|
||||
fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"lo\"}}]}\n\n")
|
||||
w.(http.Flusher).Flush()
|
||||
fmt.Fprint(w, "data: [DONE]\n\n")
|
||||
w.(http.Flusher).Flush()
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
client := newTestClient(srv.URL)
|
||||
events, errs := client.Stream(context.Background(), []model.ChatMessage{
|
||||
{Role: model.ChatRoleUser, Content: "hi"},
|
||||
}, nil)
|
||||
|
||||
var texts []string
|
||||
for ev := range events {
|
||||
if ev.Delta != "" {
|
||||
texts = append(texts, ev.Delta)
|
||||
}
|
||||
if ev.Done {
|
||||
break
|
||||
}
|
||||
}
|
||||
// 检查错误通道
|
||||
for err := range errs {
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
|
||||
assert.Equal(t, []string{"hel", "lo"}, texts)
|
||||
}
|
||||
|
||||
func TestOpenAIClient_Stream_ToolCallDeltas_MergesCorrectly(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
// 工具调用分多个 chunk 发送
|
||||
fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_1\",\"function\":{\"name\":\"search\"}}]}}]}\n\n")
|
||||
w.(http.Flusher).Flush()
|
||||
fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"{\\\"q\"}}]}}]}\n\n")
|
||||
w.(http.Flusher).Flush()
|
||||
fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"uery\\\":\\\"test\\\"}\"}}]}}]}\n\n")
|
||||
w.(http.Flusher).Flush()
|
||||
fmt.Fprint(w, "data: [DONE]\n\n")
|
||||
w.(http.Flusher).Flush()
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
client := newTestClient(srv.URL)
|
||||
events, errs := client.Stream(context.Background(), []model.ChatMessage{
|
||||
{Role: model.ChatRoleUser, Content: "hi"},
|
||||
}, nil)
|
||||
|
||||
var toolCalls []model.ChatToolCall
|
||||
for ev := range events {
|
||||
if len(ev.ToolCalls) > 0 {
|
||||
toolCalls = append(toolCalls, ev.ToolCalls...)
|
||||
}
|
||||
if ev.Done {
|
||||
break
|
||||
}
|
||||
}
|
||||
for range errs {
|
||||
}
|
||||
|
||||
assert.Len(t, toolCalls, 1)
|
||||
assert.Equal(t, "call_1", toolCalls[0].ID)
|
||||
assert.Equal(t, "search", toolCalls[0].Name)
|
||||
assert.Equal(t, `{"query":"test"}`, toolCalls[0].Arguments)
|
||||
}
|
||||
|
||||
func TestOpenAIClient_Stream_Upstream500_ReturnsError(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
fmt.Fprint(w, "error")
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
client := newTestClient(srv.URL)
|
||||
events, errs := client.Stream(context.Background(), []model.ChatMessage{
|
||||
{Role: model.ChatRoleUser, Content: "hi"},
|
||||
}, nil)
|
||||
|
||||
// 消费通道直到关闭
|
||||
var streamErr error
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
for range events {
|
||||
}
|
||||
close(done)
|
||||
}()
|
||||
for err := range errs {
|
||||
streamErr = err
|
||||
}
|
||||
<-done
|
||||
|
||||
assert.Error(t, streamErr)
|
||||
assert.Contains(t, streamErr.Error(), "500")
|
||||
}
|
||||
|
||||
// ============================================================
|
||||
// encodeMessages / encodeTools 测试
|
||||
// ============================================================
|
||||
|
||||
func TestEncodeMessages_WithToolCalls(t *testing.T) {
|
||||
msgs := []model.ChatMessage{
|
||||
{Role: model.ChatRoleAssistant, ToolCalls: []model.ChatToolCall{
|
||||
{ID: "c1", Name: "search", Arguments: `{"q":"x"}`},
|
||||
}},
|
||||
}
|
||||
encoded := encodeMessages(msgs)
|
||||
assert.Len(t, encoded, 1)
|
||||
assert.Equal(t, "assistant", encoded[0]["role"])
|
||||
calls := encoded[0]["tool_calls"].([]map[string]any)
|
||||
assert.Len(t, calls, 1)
|
||||
assert.Equal(t, "c1", calls[0]["id"])
|
||||
}
|
||||
|
||||
func TestEncodeTools_CorrectFormat(t *testing.T) {
|
||||
tools := []ToolDef{{Name: "search", Description: "search tool"}}
|
||||
encoded := encodeTools(tools)
|
||||
assert.Len(t, encoded, 1)
|
||||
assert.Equal(t, "function", encoded[0]["type"])
|
||||
fn := encoded[0]["function"].(map[string]any)
|
||||
assert.Equal(t, "search", fn["name"])
|
||||
assert.Equal(t, "search tool", fn["description"])
|
||||
}
|
||||
|
||||
func TestEncodeTools_EmptyDescription_UsesFallback(t *testing.T) {
|
||||
tools := []ToolDef{{Name: "mytool"}}
|
||||
encoded := encodeTools(tools)
|
||||
fn := encoded[0]["function"].(map[string]any)
|
||||
assert.Contains(t, fn["description"], "mytool")
|
||||
}
|
||||
|
||||
// ============================================================
|
||||
// truncate 测试
|
||||
// ============================================================
|
||||
|
||||
func TestTruncate_ShortString_Unchanged(t *testing.T) {
|
||||
assert.Equal(t, "abc", truncate("abc", 10))
|
||||
}
|
||||
|
||||
func TestTruncate_LongString_Truncated(t *testing.T) {
|
||||
result := truncate("abcdefghijk", 5)
|
||||
assert.Equal(t, "abcde...", result)
|
||||
}
|
||||
|
||||
// ============================================================
|
||||
// buildRequestBody 测试
|
||||
// ============================================================
|
||||
|
||||
func TestBuildRequestBody_Stream(t *testing.T) {
|
||||
body, err := buildRequestBody("gpt-4", []model.ChatMessage{
|
||||
{Role: model.ChatRoleUser, Content: "hi"},
|
||||
}, nil, true)
|
||||
assert.NoError(t, err)
|
||||
assert.Contains(t, string(body), `"stream":true`)
|
||||
assert.Contains(t, string(body), `"model":"gpt-4"`)
|
||||
}
|
||||
Reference in New Issue
Block a user