//go:build unit package service import ( "bytes" "context" "io" "net/http" "net/http/httptest" "strings" "testing" "time" "github.com/Wei-Shaw/sub2api/internal/config" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" ) func TestIsOpenAIImageRateLimitError(t *testing.T) { imageBody := []byte(`{"error":{"message":"Rate limit reached for gpt-image-2-codex (for limit gpt-image) in organization org on input-images per min: Limit 4000, Used 4000. Please try again in 467ms."}}`) textBody := []byte(`{"error":{"message":"Rate limit reached for gpt-5.4 in organization org on tokens per min: Limit 30000, Used 30000. Please try again in 1s."}}`) require.True(t, isOpenAIImageRateLimitError(http.StatusTooManyRequests, imageBody)) require.False(t, isOpenAIImageRateLimitError(http.StatusTooManyRequests, textBody)) require.False(t, isOpenAIImageRateLimitError(http.StatusBadRequest, imageBody)) } func TestRateLimitService_HandleOpenAIImageRateLimit_ParsesTryAgainCooldown(t *testing.T) { repo := &modelNotFoundAccountRepoStub{} svc := &RateLimitService{accountRepo: repo} account := &Account{ID: 201, Platform: PlatformOpenAI, Type: AccountTypeOAuth} body := []byte(`{"error":{"type":"rate_limit_exceeded","message":"Rate limit reached for gpt-image-2-codex (for limit gpt-image) on input-images per min. Please try again in 2s."}}`) before := time.Now() handled := svc.HandleOpenAIImageRateLimit(context.Background(), account, http.StatusTooManyRequests, http.Header{}, body) require.True(t, handled) require.Len(t, repo.modelRateLimitCalls, 1) call := repo.modelRateLimitCalls[0] require.Equal(t, account.ID, call.accountID) require.Equal(t, openAIImageGenerationRateLimitKey, call.scope) require.Equal(t, openAIImageRateLimitReason, call.reason) require.WithinDuration(t, before.Add(2*time.Second), call.resetAt, time.Second) } func TestRateLimitService_HandleOpenAIImageRateLimit_DefaultsToOneMinute(t *testing.T) { repo := &modelNotFoundAccountRepoStub{} svc := &RateLimitService{accountRepo: repo} account := &Account{ID: 202, Platform: PlatformOpenAI, Type: AccountTypeOAuth} body := []byte(`{"error":{"type":"rate_limit_exceeded","message":"Rate limit reached for gpt-image-2-codex (for limit gpt-image) on input-images per min."}}`) before := time.Now() handled := svc.HandleOpenAIImageRateLimit(context.Background(), account, http.StatusTooManyRequests, http.Header{}, body) require.True(t, handled) require.Len(t, repo.modelRateLimitCalls, 1) call := repo.modelRateLimitCalls[0] require.Equal(t, openAIImageGenerationRateLimitKey, call.scope) require.Equal(t, openAIImageRateLimitReason, call.reason) require.WithinDuration(t, before.Add(time.Minute), call.resetAt, time.Second) } func TestOpenAIGatewayService_HandleOpenAIAccountUpstreamError_ImageRateLimitDoesNotBlockWholeAccount(t *testing.T) { repo := &modelNotFoundAccountRepoStub{} svc := &OpenAIGatewayService{rateLimitService: &RateLimitService{accountRepo: repo}} account := &Account{ID: 203, Platform: PlatformOpenAI, Type: AccountTypeOAuth} body := []byte(`{"error":{"type":"rate_limit_exceeded","message":"Rate limit reached for gpt-image-2-codex (for limit gpt-image) on input-images per min. Please try again in 1s."}}`) disabled := svc.handleOpenAIAccountUpstreamError(context.Background(), account, http.StatusTooManyRequests, http.Header{}, body, "gpt-image-2") require.False(t, disabled) require.Len(t, repo.modelRateLimitCalls, 1) require.Equal(t, openAIImageGenerationRateLimitKey, repo.modelRateLimitCalls[0].scope) _, wholeAccountBlocked := svc.openaiAccountRuntimeBlockUntil.Load(account.ID) require.False(t, wholeAccountBlocked) } func TestOpenAIGatewayServiceForwardImages_ImageRateLimitReturnsFailoverAndCoolsCapability(t *testing.T) { gin.SetMode(gin.TestMode) repo := &modelNotFoundAccountRepoStub{} body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat"}`) errorBody := `{"error":{"type":"rate_limit_exceeded","message":"Rate limit reached for gpt-image-2-codex (for limit gpt-image) in organization org on input-images per min: Limit 4000, Used 4000. Please try again in 1s."}}` req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)) req.Header.Set("Content-Type", "application/json") rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) c.Request = req svc := &OpenAIGatewayService{ rateLimitService: &RateLimitService{accountRepo: repo}, httpUpstream: &httpUpstreamRecorder{ resp: &http.Response{ StatusCode: http.StatusTooManyRequests, Header: http.Header{"X-Request-Id": []string{"req_img_rate_limited"}}, Body: io.NopCloser(strings.NewReader(errorBody)), }, }, } parsed, err := svc.ParseOpenAIImagesRequest(c, body) require.NoError(t, err) account := &Account{ ID: 204, Name: "openai-oauth", Platform: PlatformOpenAI, Type: AccountTypeOAuth, Credentials: map[string]any{ "access_token": "token-123", }, } result, err := svc.ForwardImages(context.Background(), c, account, body, parsed, "") require.Nil(t, result) var failoverErr *UpstreamFailoverError require.ErrorAs(t, err, &failoverErr) require.Equal(t, http.StatusTooManyRequests, failoverErr.StatusCode) require.Contains(t, string(failoverErr.ResponseBody), "input-images per min") require.Len(t, repo.modelRateLimitCalls, 1) require.Equal(t, openAIImageGenerationRateLimitKey, repo.modelRateLimitCalls[0].scope) } // issue #6171:上游"回文字没回图"是**这一轮**的结果(模型选择了说话),不是账号能力 // 失效。它同时被判为可重试(502)并驱动 failover,若还写 30 分钟账号级冷却,一次闲聊 // 回复就会沿号池把每个被重试到的账号依次冷却掉。冷却仍保留给结构化上游证据,见 // TestOpenAIGatewayServiceForwardImages_StructuredUnavailableCoolsImageCapability。 func TestOpenAIGatewayServiceForwardImages_TextFallbackDoesNotCoolImageCapability(t *testing.T) { gin.SetMode(gin.TestMode) repo := &modelNotFoundAccountRepoStub{} body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat"}`) upstreamSSE := "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"r\",\"status\":\"completed\",\"model\":\"gpt-5.4-mini\",\"output\":[{\"type\":\"message\",\"content\":[{\"type\":\"output_text\",\"text\":\"Here's a polished image prompt for your request.\"}]}]}}\n\n" req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)) req.Header.Set("Content-Type", "application/json") rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) c.Request = req svc := &OpenAIGatewayService{ accountRepo: repo, httpUpstream: &httpUpstreamRecorder{ resp: &http.Response{ StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"text/event-stream"}}, Body: io.NopCloser(strings.NewReader(upstreamSSE)), }, }, } parsed, err := svc.ParseOpenAIImagesRequest(c, body) require.NoError(t, err) account := &Account{ ID: 205, Name: "openai-oauth", Platform: PlatformOpenAI, Type: AccountTypeOAuth, Credentials: map[string]any{ "access_token": "token-123", }, } result, err := svc.ForwardImages(context.Background(), c, account, body, parsed, "") require.Nil(t, result) require.Error(t, err) var failoverErr *UpstreamFailoverError require.ErrorAs(t, err, &failoverErr) require.False(t, failoverErr.RetryableOnSameAccount) // 换号行为不变:该判据仍足以放弃本账号重试这一次请求…… require.Equal(t, http.StatusBadGateway, failoverErr.StatusCode) // ……但不再写任何账号级状态,否则重试会把冷却一路刷到整个号池。 require.Empty(t, repo.modelRateLimitCalls, "模型回文字只说明这一轮没出图,不构成账号 30 分钟不可用的证据") } // 对照不变式:上游 error 帧点名 image_generation_unavailable 时仍写冷却, // 保证 #6171 的修复没有把这项能力保护整个废掉。 func TestOpenAIGatewayServiceForwardImages_StructuredUnavailableCoolsImageCapability(t *testing.T) { gin.SetMode(gin.TestMode) repo := &modelNotFoundAccountRepoStub{} body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat"}`) upstreamSSE := "data: {\"type\":\"response.failed\",\"response\":{\"id\":\"r\",\"error\":" + "{\"type\":\"upstream_error\",\"code\":\"image_generation_unavailable\"," + "\"message\":\"image generation tool is not available for this account\"}}}\n\n" req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)) req.Header.Set("Content-Type", "application/json") rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) c.Request = req svc := &OpenAIGatewayService{ accountRepo: repo, httpUpstream: &httpUpstreamRecorder{ resp: &http.Response{ StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"text/event-stream"}}, Body: io.NopCloser(strings.NewReader(upstreamSSE)), }, }, } parsed, err := svc.ParseOpenAIImagesRequest(c, body) require.NoError(t, err) account := &Account{ ID: 206, Name: "openai-oauth", Platform: PlatformOpenAI, Type: AccountTypeOAuth, Credentials: map[string]any{ "access_token": "token-123", }, } before := time.Now() result, err := svc.ForwardImages(context.Background(), c, account, body, parsed, "") require.Nil(t, result) require.Error(t, err) require.Len(t, repo.modelRateLimitCalls, 1) call := repo.modelRateLimitCalls[0] require.Equal(t, account.ID, call.accountID) require.Equal(t, openAIImageGenerationRateLimitKey, call.scope) require.Equal(t, openAIImagesOAuthUnavailableReason, call.reason) require.WithinDuration(t, before.Add(openAIImagesOAuthUnavailableDefaultCooldown), call.resetAt, time.Second) } func TestOpenAIGatewayService_CoolOpenAIImagesOAuthToolUsesConfiguredCooldown(t *testing.T) { accountRepo := &modelNotFoundAccountRepoStub{} settingRepo := newMockSettingRepo() settingRepo.data[SettingKeyOpenAIImagesOAuthUnavailableCooldownSettings] = `{"cooldown_minutes":7}` svc := &OpenAIGatewayService{ accountRepo: accountRepo, settingService: NewSettingService(settingRepo, &config.Config{}), } before := time.Now() svc.coolOpenAIImagesOAuthTool(context.Background(), &Account{ID: 206, Platform: PlatformOpenAI, Type: AccountTypeOAuth}) require.Len(t, accountRepo.modelRateLimitCalls, 1) require.WithinDuration(t, before.Add(7*time.Minute), accountRepo.modelRateLimitCalls[0].resetAt, time.Second) } func TestOpenAIGatewayServiceForwardImages_CapabilityLossCoolsImageScope(t *testing.T) { gin.SetMode(gin.TestMode) repo := &modelNotFoundAccountRepoStub{} body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat"}`) errorBody := `{"error":{"message":"Tool choice 'image_generation' not found in 'tools' parameter.","param":"tool_choice","type":"invalid_request_error"}}` req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)) req.Header.Set("Content-Type", "application/json") rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) c.Request = req svc := &OpenAIGatewayService{ rateLimitService: &RateLimitService{accountRepo: repo}, httpUpstream: &httpUpstreamRecorder{ resp: &http.Response{ StatusCode: http.StatusBadRequest, Header: http.Header{"X-Request-Id": []string{"req_img_capability_lost"}}, Body: io.NopCloser(strings.NewReader(errorBody)), }, }, } parsed, err := svc.ParseOpenAIImagesRequest(c, body) require.NoError(t, err) account := &Account{ ID: 205, Name: "openai-oauth", Platform: PlatformOpenAI, Type: AccountTypeOAuth, Credentials: map[string]any{ "access_token": "token-123", }, } before := time.Now() result, err := svc.ForwardImages(context.Background(), c, account, body, parsed, "") require.Nil(t, result) require.Error(t, err) require.Len(t, repo.modelRateLimitCalls, 1) call := repo.modelRateLimitCalls[0] require.Equal(t, account.ID, call.accountID) require.Equal(t, openAIImageGenerationRateLimitKey, call.scope) require.Equal(t, openAIImageCapabilityLossReason, call.reason) require.WithinDuration(t, before.Add(openAIImageCapabilityLossCooldown), call.resetAt, time.Second) } func TestOpenAIGatewayServiceHandleUpstreamError_PassthroughCapabilityLossDoesNotCool(t *testing.T) { repo := &modelNotFoundAccountRepoStub{} svc := &OpenAIGatewayService{rateLimitService: &RateLimitService{accountRepo: repo}} account := &Account{ID: 206, Platform: PlatformOpenAI, Type: AccountTypeOAuth} body := []byte(`{"error":{"message":"Tool choice 'image_generation' not found in 'tools' parameter.","param":"tool_choice","type":"invalid_request_error"}}`) disabled := svc.handleOpenAIAccountUpstreamError(context.Background(), account, http.StatusBadRequest, http.Header{}, body, "gpt-5.5") require.False(t, disabled) require.Empty(t, repo.modelRateLimitCalls) _, wholeAccountBlocked := svc.openaiAccountRuntimeBlockUntil.Load(account.ID) require.False(t, wholeAccountBlocked) } func TestRateLimitServiceHandleOpenAIImageCapabilityLoss_IgnoresGenericBadRequest(t *testing.T) { repo := &modelNotFoundAccountRepoStub{} svc := &RateLimitService{accountRepo: repo} account := &Account{ID: 207, Platform: PlatformOpenAI, Type: AccountTypeOAuth} body := []byte(`{"error":{"message":"Invalid type for input[0].arguments"}}`) handled := svc.HandleOpenAIImageCapabilityLoss(context.Background(), account, http.StatusBadRequest, body) require.False(t, handled) require.Empty(t, repo.modelRateLimitCalls) } func TestRateLimitServiceHandleOpenAIImageCapabilityLoss_RespectsPlatformAndErrorCodePolicy(t *testing.T) { body := []byte(`{"error":{"message":"Tool choice 'image_generation' not found in 'tools' parameter.","param":"tool_choice","type":"invalid_request_error"}}`) t.Run("non_openai_platform", func(t *testing.T) { repo := &modelNotFoundAccountRepoStub{} svc := &RateLimitService{accountRepo: repo} account := &Account{ID: 208, Platform: PlatformAnthropic, Type: AccountTypeOAuth} handled := svc.HandleOpenAIImageCapabilityLoss(context.Background(), account, http.StatusBadRequest, body) require.False(t, handled) require.Empty(t, repo.modelRateLimitCalls) }) t.Run("custom_error_code_policy_excludes_400", func(t *testing.T) { repo := &modelNotFoundAccountRepoStub{} svc := &RateLimitService{accountRepo: repo} account := &Account{ ID: 209, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Credentials: map[string]any{ "custom_error_codes_enabled": true, "custom_error_codes": []any{float64(http.StatusTooManyRequests)}, }, } require.False(t, account.ShouldHandleErrorCode(http.StatusBadRequest)) handled := svc.HandleOpenAIImageCapabilityLoss(context.Background(), account, http.StatusBadRequest, body) require.False(t, handled) require.Empty(t, repo.modelRateLimitCalls) }) } func TestIsOpenAIImageCapabilityLossError(t *testing.T) { capabilityLossBody := []byte(`{"error":{"message":"Tool choice 'image_generation' not found in 'tools' parameter.","param":"tool_choice","type":"invalid_request_error"}}`) genericBadRequestBody := []byte(`{"error":{"message":"Invalid type for input[0].arguments"}}`) require.True(t, isOpenAIImageCapabilityLossError(http.StatusBadRequest, capabilityLossBody)) require.False(t, isOpenAIImageCapabilityLossError(http.StatusBadRequest, genericBadRequestBody)) require.False(t, isOpenAIImageCapabilityLossError(http.StatusTooManyRequests, capabilityLossBody)) }