252 lines
6.8 KiB
Go
252 lines
6.8 KiB
Go
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)
|
||
}
|
||
})
|
||
}
|
||
}
|