package llm import ( "context" "fmt" "net/http" "net/http/httptest" "strings" "testing" ) // 用一个 OpenAI 兼容的 mock 服务模拟本地 vLLM/Ollama,验证: // 1) 本地 provider 的流式响应能被 TransformOpenAIStream 正确解析; // 2) 未配置密钥时不发送 Authorization 头(本地无鉴权端点)。 func TestLocalProvider_StreamingAndNoAuthHeader(t *testing.T) { var gotAuth string var sawAuthHeader bool srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { gotAuth = r.Header.Get("Authorization") _, sawAuthHeader = r.Header["Authorization"] w.Header().Set("Content-Type", "text/event-stream") flusher, _ := w.(http.Flusher) for _, chunk := range []string{ `{"id":"cmpl-1","model":"local-model","choices":[{"delta":{"content":"你好"}}]}`, `{"id":"cmpl-1","model":"local-model","choices":[{"delta":{"content":",世界"}}]}`, } { fmt.Fprintf(w, "data: %s\n\n", chunk) if flusher != nil { flusher.Flush() } } fmt.Fprint(w, "data: [DONE]\n\n") if flusher != nil { flusher.Flush() } })) defer srv.Close() mgr := NewManager() // 密钥留空,模拟本地无鉴权端点 mgr.Register("local", NewOpenAIProvider("", srv.URL, "local-model")) body, err := mgr.ChatStream(context.Background(), "local", &ChatRequest{ Messages: []Message{{Role: RoleUser, Content: "hi"}}, }) if err != nil { t.Fatalf("ChatStream 出错: %v", err) } defer body.Close() var sb strings.Builder var ended bool if err := TransformOpenAIStream(body, func(ev StreamEvent) { if ev.Answer != "" { sb.WriteString(ev.Answer) } if ev.Event == "message_end" { ended = true } }); err != nil { t.Fatalf("解析流出错: %v", err) } if sb.String() != "你好,世界" { t.Fatalf("流式拼接结果不符: %q", sb.String()) } if !ended { t.Fatal("未收到 message_end 事件") } if sawAuthHeader || gotAuth != "" { t.Fatalf("空密钥时不应发送 Authorization 头,实际: %q", gotAuth) } } // 配置了密钥时应发送 Authorization 头(云端/带鉴权的本地服务)。 func TestOpenAIProvider_SendsAuthHeaderWhenKeySet(t *testing.T) { var gotAuth string srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { gotAuth = r.Header.Get("Authorization") w.Header().Set("Content-Type", "application/json") fmt.Fprint(w, `{"id":"1","model":"m","choices":[{"message":{"content":"ok"}}],"usage":{"total_tokens":3}}`) })) defer srv.Close() p := NewOpenAIProvider("test-key", srv.URL, "m") resp, err := p.ChatCompletion(context.Background(), &ChatRequest{ Messages: []Message{{Role: RoleUser, Content: "hi"}}, }) if err != nil { t.Fatalf("ChatCompletion 出错: %v", err) } if resp.Content != "ok" { t.Fatalf("响应内容不符: %q", resp.Content) } if gotAuth != "Bearer test-key" { t.Fatalf("应发送 Bearer 密钥头,实际: %q", gotAuth) } } // 未注册的 provider 名应回退到 fallback。 func TestManager_FallbackResolution(t *testing.T) { mgr := NewManager() mgr.Register("local", NewOpenAIProvider("", "http://localhost:9", "m")) mgr.SetFallback("local") if _, err := mgr.GetProvider("does-not-exist"); err != nil { t.Fatalf("未知 provider 应回退到 fallback,却报错: %v", err) } }