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("bad gateway")) 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("连续越界应返回错误") } }