test: 编写 LLM 测试（mock SSE 流，覆盖成功/图片/历史/API 错误/超时/Usage）
This commit is contained in:
251
backend/internal/ai/llm/openai_test.go
Normal file
251
backend/internal/ai/llm/openai_test.go
Normal file
@@ -0,0 +1,251 @@
|
|||||||
|
package llm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"go.uber.org/zap"
|
||||||
|
|
||||||
|
"github.com/hhs/camtalk/internal/models"
|
||||||
|
)
|
||||||
|
|
||||||
|
// mockLLMServer 创建模拟 OpenAI SSE 流式响应的 HTTP 服务器。
|
||||||
|
func mockLLMServer(t *testing.T, handler http.HandlerFunc) *httptest.Server {
|
||||||
|
t.Helper()
|
||||||
|
return httptest.NewServer(handler)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestOpenAIService_ChatStream_Success(t *testing.T) {
|
||||||
|
srv := mockLLMServer(t, func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
// 验证请求
|
||||||
|
if r.Method != http.MethodPost {
|
||||||
|
t.Errorf("method = %s, want POST", r.Method)
|
||||||
|
}
|
||||||
|
if !strings.Contains(r.URL.Path, "/chat/completions") {
|
||||||
|
t.Errorf("path = %s, should contain /chat/completions", r.URL.Path)
|
||||||
|
}
|
||||||
|
auth := r.Header.Get("Authorization")
|
||||||
|
if auth != "Bearer test-key" {
|
||||||
|
t.Errorf("Authorization = %q, want %q", auth, "Bearer test-key")
|
||||||
|
}
|
||||||
|
|
||||||
|
w.Header().Set("Content-Type", "text/event-stream")
|
||||||
|
flusher, ok := w.(http.Flusher)
|
||||||
|
if !ok {
|
||||||
|
t.Fatal("ResponseWriter does not support Flusher")
|
||||||
|
}
|
||||||
|
|
||||||
|
// 发送几个 delta
|
||||||
|
deltas := []string{"你好", "世界", "！"}
|
||||||
|
for _, d := range deltas {
|
||||||
|
fmt.Fprintf(w, "data: {\"choices\":[{\"delta\":{\"content\":\"%s\"}}],\"model\":\"gpt-4o\"}\n\n", d)
|
||||||
|
flusher.Flush()
|
||||||
|
}
|
||||||
|
|
||||||
|
// 发送 [DONE]
|
||||||
|
fmt.Fprintf(w, "data: [DONE]\n\n")
|
||||||
|
flusher.Flush()
|
||||||
|
})
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
svc := NewOpenAIService("test-key", "gpt-4o", srv.URL, 10, zap.NewNop().Sugar())
|
||||||
|
|
||||||
|
ch, err := svc.ChatStream(context.Background(), Request{
|
||||||
|
Text: "这是什么？",
|
||||||
|
Language: "zh-CN",
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ChatStream() error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
var chunks []Chunk
|
||||||
|
for c := range ch {
|
||||||
|
chunks = append(chunks, c)
|
||||||
|
}
|
||||||
|
|
||||||
|
// 应该有 3 个文本 chunk + 1 个 Done chunk
|
||||||
|
if len(chunks) != 4 {
|
||||||
|
t.Fatalf("got %d chunks, want 4", len(chunks))
|
||||||
|
}
|
||||||
|
|
||||||
|
// 验证文本内容
|
||||||
|
if chunks[0].Delta != "你好" {
|
||||||
|
t.Errorf("chunk[0].Delta = %q, want %q", chunks[0].Delta, "你好")
|
||||||
|
}
|
||||||
|
if chunks[1].Delta != "世界" {
|
||||||
|
t.Errorf("chunk[1].Delta = %q, want %q", chunks[1].Delta, "世界")
|
||||||
|
}
|
||||||
|
|
||||||
|
// 验证最后一个 chunk 是 Done
|
||||||
|
last := chunks[len(chunks)-1]
|
||||||
|
if !last.Done {
|
||||||
|
t.Error("last chunk should be Done")
|
||||||
|
}
|
||||||
|
if last.Model != "gpt-4o" {
|
||||||
|
t.Errorf("last chunk Model = %q, want %q", last.Model, "gpt-4o")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestOpenAIService_ChatStream_WithImage(t *testing.T) {
|
||||||
|
srv := mockLLMServer(t, func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "text/event-stream")
|
||||||
|
fmt.Fprintf(w, "data: {\"choices\":[{\"delta\":{\"content\":\"ok\"}}],\"model\":\"gpt-4o\"}\n\n")
|
||||||
|
fmt.Fprintf(w, "data: [DONE]\n\n")
|
||||||
|
})
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
svc := NewOpenAIService("test-key", "gpt-4o", srv.URL, 10, zap.NewNop().Sugar())
|
||||||
|
|
||||||
|
ch, err := svc.ChatStream(context.Background(), Request{
|
||||||
|
Image: []byte("fake-jpeg-data"),
|
||||||
|
Text: "描述图片",
|
||||||
|
Language: "zh-CN",
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ChatStream() error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// 消费 channel
|
||||||
|
for range ch {
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestOpenAIService_ChatStream_WithHistory(t *testing.T) {
|
||||||
|
srv := mockLLMServer(t, func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "text/event-stream")
|
||||||
|
fmt.Fprintf(w, "data: {\"choices\":[{\"delta\":{\"content\":\"ok\"}}],\"model\":\"gpt-4o\"}\n\n")
|
||||||
|
fmt.Fprintf(w, "data: [DONE]\n\n")
|
||||||
|
})
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
svc := NewOpenAIService("test-key", "gpt-4o", srv.URL, 10, zap.NewNop().Sugar())
|
||||||
|
|
||||||
|
ch, err := svc.ChatStream(context.Background(), Request{
|
||||||
|
Text: "继续",
|
||||||
|
Language: "zh-CN",
|
||||||
|
History: []models.Message{
|
||||||
|
{Role: "user", Content: "你好"},
|
||||||
|
{Role: "assistant", Content: "你好！有什么可以帮助你的吗？"},
|
||||||
|
},
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ChatStream() error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
for range ch {
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestOpenAIService_ChatStream_APIError(t *testing.T) {
|
||||||
|
srv := mockLLMServer(t, func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
w.WriteHeader(http.StatusUnauthorized)
|
||||||
|
fmt.Fprintf(w, `{"error":{"message":"Invalid API key"}}`)
|
||||||
|
})
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
svc := NewOpenAIService("bad-key", "gpt-4o", srv.URL, 10, zap.NewNop().Sugar())
|
||||||
|
|
||||||
|
_, err := svc.ChatStream(context.Background(), Request{
|
||||||
|
Text: "test",
|
||||||
|
})
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("ChatStream() should return error for 401")
|
||||||
|
}
|
||||||
|
if !strings.Contains(err.Error(), "401") {
|
||||||
|
t.Errorf("error should mention 401, got: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestOpenAIService_ChatStream_Timeout(t *testing.T) {
|
||||||
|
srv := mockLLMServer(t, func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
// 模拟慢响应
|
||||||
|
time.Sleep(5 * time.Second)
|
||||||
|
w.Header().Set("Content-Type", "text/event-stream")
|
||||||
|
fmt.Fprintf(w, "data: {\"choices\":[{\"delta\":{\"content\":\"late\"}}]}\n\n")
|
||||||
|
fmt.Fprintf(w, "data: [DONE]\n\n")
|
||||||
|
})
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
svc := NewOpenAIService("test-key", "gpt-4o", srv.URL, 1, zap.NewNop().Sugar()) // 1s timeout
|
||||||
|
|
||||||
|
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||||
|
defer cancel()
|
||||||
|
|
||||||
|
ch, err := svc.ChatStream(ctx, Request{Text: "test"})
|
||||||
|
if err != nil {
|
||||||
|
// 超时可能在建立连接时或读取时发生
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// 如果连接成功，消费 channel 应该超时
|
||||||
|
var gotContent bool
|
||||||
|
for c := range ch {
|
||||||
|
if c.Delta != "" {
|
||||||
|
gotContent = true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if gotContent {
|
||||||
|
t.Error("should not receive content before timeout")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestOpenAIService_ChatStream_UsageInResponse(t *testing.T) {
|
||||||
|
srv := mockLLMServer(t, func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "text/event-stream")
|
||||||
|
// 带 usage 的最后一个 chunk
|
||||||
|
fmt.Fprintf(w, "data: {\"choices\":[{\"delta\":{\"content\":\"hi\"},\"finish_reason\":\"stop\"}],\"model\":\"gpt-4o\",\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":5,\"total_tokens\":15}}\n\n")
|
||||||
|
fmt.Fprintf(w, "data: [DONE]\n\n")
|
||||||
|
})
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
svc := NewOpenAIService("test-key", "gpt-4o", srv.URL, 10, zap.NewNop().Sugar())
|
||||||
|
|
||||||
|
ch, err := svc.ChatStream(context.Background(), Request{Text: "test"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ChatStream() error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
var last Chunk
|
||||||
|
for c := range ch {
|
||||||
|
last = c
|
||||||
|
}
|
||||||
|
|
||||||
|
if !last.Done {
|
||||||
|
t.Error("last chunk should be Done")
|
||||||
|
}
|
||||||
|
if last.TokensUsed == nil {
|
||||||
|
t.Fatal("last chunk should have TokensUsed")
|
||||||
|
}
|
||||||
|
if last.TokensUsed.Total != 15 {
|
||||||
|
t.Errorf("TokensUsed.Total = %d, want 15", last.TokensUsed.Total)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBuildSystemPrompt(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
language string
|
||||||
|
detailLevel string
|
||||||
|
wantContain string
|
||||||
|
}{
|
||||||
|
{"chinese default", "zh-CN", "", "视觉助手"},
|
||||||
|
{"chinese high", "zh-CN", "high", "更详细"},
|
||||||
|
{"english default", "en", "", "visual assistant"},
|
||||||
|
{"english high", "en", "high", "detailed"},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
got := BuildSystemPrompt(tt.language, tt.detailLevel)
|
||||||
|
if !strings.Contains(got, tt.wantContain) {
|
||||||
|
t.Errorf("BuildSystemPrompt(%q, %q) should contain %q", tt.language, tt.detailLevel, tt.wantContain)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user