172 lines
5.5 KiB
Go
172 lines
5.5 KiB
Go
package ai
|
||
|
||
import (
|
||
"context"
|
||
"fmt"
|
||
"io"
|
||
"net/http"
|
||
"net/http/httptest"
|
||
"strings"
|
||
"sync/atomic"
|
||
"testing"
|
||
"time"
|
||
)
|
||
|
||
// newTestClient 构造指向本地假服务器的 LLM 客户端
|
||
func newTestClient(srv *httptest.Server) *LLMClient {
|
||
return &LLMClient{
|
||
BaseURL: srv.URL,
|
||
APIKey: "test-key",
|
||
Model: "test-model",
|
||
http: srv.Client(),
|
||
}
|
||
}
|
||
|
||
// okBody 组装一条合法的 OpenAI 格式响应
|
||
func okBody(content string) string {
|
||
return fmt.Sprintf(`{"choices":[{"message":{"content":%q}}]}`, content)
|
||
}
|
||
|
||
// TestChatOneValidWins 三路并发中只要有一路合格就应成功（其余 500/非法 JSON）
|
||
func TestChatOneValidWins(t *testing.T) {
|
||
var calls int32
|
||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
switch atomic.AddInt32(&calls, 1) {
|
||
case 1:
|
||
w.WriteHeader(500)
|
||
w.Write([]byte("overloaded"))
|
||
case 2:
|
||
w.Write([]byte("<html>bad gateway</html>"))
|
||
default:
|
||
w.Write([]byte(okBody("你好")))
|
||
}
|
||
}))
|
||
defer srv.Close()
|
||
out, err := newTestClient(srv).Chat(context.Background(), "s", "u", 0.5)
|
||
if err != nil || out != "你好" {
|
||
t.Fatalf("有一路合格就应成功，got out=%q err=%v", out, err)
|
||
}
|
||
}
|
||
|
||
// TestChatAllFailAggregates 三路全部失败：应发满 3 路并汇总去重后的错误
|
||
func TestChatAllFailAggregates(t *testing.T) {
|
||
var calls int32
|
||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
atomic.AddInt32(&calls, 1)
|
||
w.WriteHeader(401)
|
||
w.Write([]byte(`{"error":{"message":"invalid api key"}}`))
|
||
}))
|
||
defer srv.Close()
|
||
_, err := newTestClient(srv).Chat(context.Background(), "s", "u", 0.5)
|
||
if err == nil {
|
||
t.Fatal("全部失败应返回错误")
|
||
}
|
||
if got := atomic.LoadInt32(&calls); got != chatParallel {
|
||
t.Fatalf("应并发发出 %d 路请求，实际 %d", chatParallel, got)
|
||
}
|
||
// 相同错误应去重，不会把同一句话重复三遍
|
||
if strings.Count(err.Error(), "HTTP 401") != 1 {
|
||
t.Fatalf("相同错误应去重，实际：%v", err)
|
||
}
|
||
}
|
||
|
||
// TestChatEarlyReturn 一路秒回合格结果时不应等慢的两路跑完
|
||
func TestChatEarlyReturn(t *testing.T) {
|
||
var calls int32
|
||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
if atomic.AddInt32(&calls, 1) == 1 {
|
||
w.Write([]byte(okBody("快答")))
|
||
return
|
||
}
|
||
// 其余两路拖 3 秒（可被取消提前结束）
|
||
select {
|
||
case <-r.Context().Done():
|
||
case <-time.After(3 * time.Second):
|
||
}
|
||
w.Write([]byte(okBody("慢答")))
|
||
}))
|
||
defer srv.Close()
|
||
start := time.Now()
|
||
out, err := newTestClient(srv).Chat(context.Background(), "s", "u", 0.5)
|
||
if err != nil || out != "快答" {
|
||
t.Fatalf("应采用最快的合格结果，got out=%q err=%v", out, err)
|
||
}
|
||
if elapsed := time.Since(start); elapsed > 1500*time.Millisecond {
|
||
t.Fatalf("拿到合格结果后应立即返回，实际耗时 %v", elapsed)
|
||
}
|
||
}
|
||
|
||
// TestChatValidationChain 空 choices、空内容都应被校验拦下，仅内容合格的一路胜出
|
||
func TestChatValidationChain(t *testing.T) {
|
||
var calls int32
|
||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
switch atomic.AddInt32(&calls, 1) {
|
||
case 1:
|
||
w.Write([]byte(`{"choices":[]}`))
|
||
case 2:
|
||
w.Write([]byte(okBody(" ")))
|
||
default:
|
||
w.Write([]byte(okBody("有效内容")))
|
||
}
|
||
}))
|
||
defer srv.Close()
|
||
out, err := newTestClient(srv).Chat(context.Background(), "s", "u", 0.5)
|
||
if err != nil || out != "有效内容" {
|
||
t.Fatalf("校验链应过滤空返回，got out=%q err=%v", out, err)
|
||
}
|
||
}
|
||
|
||
// TestChatRespectContext 调用方超时应尽快中止等待
|
||
func TestChatRespectContext(t *testing.T) {
|
||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
select {
|
||
case <-r.Context().Done():
|
||
case <-time.After(2 * time.Second):
|
||
}
|
||
w.WriteHeader(500)
|
||
}))
|
||
defer srv.Close()
|
||
ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
|
||
defer cancel()
|
||
start := time.Now()
|
||
_, err := newTestClient(srv).Chat(ctx, "s", "u", 0.5)
|
||
if err == nil {
|
||
t.Fatal("超时应返回错误")
|
||
}
|
||
if elapsed := time.Since(start); elapsed > time.Second {
|
||
t.Fatalf("超时后应立即中止，实际耗时 %v", elapsed)
|
||
}
|
||
}
|
||
|
||
// TestChatDecisionReask 输出不是决策 JSON 时应追加纠错提示重问一次
|
||
// （用请求体是否含纠错标记做确定性分流，避免并发时序影响）
|
||
func TestChatDecisionReask(t *testing.T) {
|
||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
body, _ := io.ReadAll(r.Body)
|
||
if strings.Contains(string(body), "无法解析") {
|
||
// 第二轮（带纠错提示）：返回合法决策
|
||
w.Write([]byte(okBody(`{"choice":1,"say":"就选它"}`)))
|
||
return
|
||
}
|
||
// 第一轮：返回闲聊文本（通过 Chat 校验但不是决策 JSON）
|
||
w.Write([]byte(okBody("我出王炸！哈哈哈")))
|
||
}))
|
||
defer srv.Close()
|
||
d, err := newTestClient(srv).ChatDecision(context.Background(), "s", "u", 0.5, 3)
|
||
if err != nil || d.Choice != 1 {
|
||
t.Fatalf("纠错重问应成功，got d=%+v err=%v", d, err)
|
||
}
|
||
}
|
||
|
||
// TestChatDecisionGiveUp choice 连续越界：两轮后放弃报错（调用方走规则兜底）
|
||
func TestChatDecisionGiveUp(t *testing.T) {
|
||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
w.Write([]byte(okBody(`{"choice":99,"say":"乱选"}`)))
|
||
}))
|
||
defer srv.Close()
|
||
_, err := newTestClient(srv).ChatDecision(context.Background(), "s", "u", 0.5, 3)
|
||
if err == nil {
|
||
t.Fatal("连续越界应返回错误")
|
||
}
|
||
}
|