package llm import ( "context" "encoding/json" "net/http" "net/http/httptest" "strings" "testing" "time" ) // TestRetryOn500:前两次 500,第三次成功——provider 应自动重试。 func TestRetryOn500(t *testing.T) { calls := 0 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { calls++ if calls <= 2 { w.WriteHeader(http.StatusInternalServerError) return } w.Header().Set("Content-Type", "text/event-stream") w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"ok\"}}]}\n\n")) w.Write([]byte("data: [DONE]\n\n")) })) defer srv.Close() p := &OpenAIProvider{BaseURL: srv.URL, APIKey: "k", Model: "m", HTTP: srv.Client()} final, err := p.Stream(context.Background(), "sys", []Message{{Role: RoleUser, Content: "hi"}}, nil, StreamOptions{}, nil) if err != nil { t.Fatalf("stream should succeed after retries: %v", err) } if final.Content != "ok" || calls != 3 { t.Fatalf("content=%q calls=%d", final.Content, calls) } } // TestNoRetryOn400:4xx 属不可重试错误,应立即返回。 func TestNoRetryOn400(t *testing.T) { calls := 0 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { calls++ w.WriteHeader(http.StatusBadRequest) w.Write([]byte(`{"error":"bad request"}`)) })) defer srv.Close() p := &OpenAIProvider{BaseURL: srv.URL, APIKey: "k", Model: "m", HTTP: srv.Client()} _, err := p.Stream(context.Background(), "sys", nil, nil, StreamOptions{}, nil) if err == nil || !strings.Contains(err.Error(), "400") { t.Fatalf("expected 400 error, got %v", err) } if calls != 1 { t.Fatalf("4xx must not retry, calls=%d", calls) } } // TestAnthropicCacheControlAndMaxTokens:system/tools 标记 cache_control, // max_tokens 来自 provider 配置。 func TestAnthropicCacheControlAndMaxTokens(t *testing.T) { var body map[string]any srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _ = json.NewDecoder(r.Body).Decode(&body) w.Header().Set("Content-Type", "text/event-stream") // 带一句正文:只发 message_stop 的空流现在会被当成失败(见 ErrEmptyCompletion), // 而这条用例断言的是**请求体**形状。 w.Write([]byte("data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"ok\"}}\n\n")) w.Write([]byte("data: {\"type\":\"message_stop\"}\n\n")) })) defer srv.Close() p := &AnthropicProvider{BaseURL: srv.URL, APIKey: "k", Model: "m", MaxTokens: 4096, HTTP: srv.Client()} _, err := p.Stream(context.Background(), "system prompt", []Message{{Role: RoleUser, Content: "hi"}}, []ToolDef{{Name: "read", Description: "d", Parameters: map[string]any{"type": "object"}}}, StreamOptions{}, nil) if err != nil { t.Fatal(err) } if body["max_tokens"].(float64) != 4096 { t.Fatalf("max_tokens: %v", body["max_tokens"]) } sys := body["system"].([]any) s0 := sys[0].(map[string]any) if s0["cache_control"].(map[string]any)["type"] != "ephemeral" { t.Fatalf("system cache_control missing: %v", s0) } tools := body["tools"].([]any) last := tools[len(tools)-1].(map[string]any) if last["cache_control"].(map[string]any)["type"] != "ephemeral" { t.Fatalf("tools cache_control missing: %v", last) } } func TestParseRetryAfter(t *testing.T) { if got := parseRetryAfter("3"); got != 3*time.Second { t.Fatalf("got %v", got) } if got := parseRetryAfter(""); got != 0 { t.Fatalf("got %v", got) } if got := parseRetryAfter("99999"); got != 0 { // 上限保护 t.Fatalf("got %v", got) } }