from __future__ import annotations import os import tempfile import unittest from datetime import datetime, timedelta, timezone from pathlib import Path from types import SimpleNamespace from unittest.mock import MagicMock, patch os.environ.setdefault("CHATGPT2API_AUTH_KEY", "test-auth") from services.account_service import AccountService from services.config import config from services.openai_backend_api import OpenAIBackendAPI from services.protocol import conversation from services.storage.json_storage import JSONStorageBackend from services.upload_limits import UploadLimit, UploadUnavailableError, parse_upload_limit from utils.helper import UpstreamHTTPError, ensure_ok def limited(message="你已达到文件上传上限。请20小时 内重试。", **kwargs): return UpstreamHTTPError("/backend-api/files", 429, { "detail": {"code": "throttled", "message": message}, }, **kwargs) class UploadLimitParsingTests(unittest.TestCase): def test_chinese_and_english_reset_times(self): now = datetime(2026, 9, 17, tzinfo=timezone.utc) for message, seconds in [ ("你已达到文件上传上限。请20小时 内重试。", 72000), ("Upload limit reached. Try again in 1 hour 30 minutes.", 5400), ("请30分钟后重试", 1800), ("Try again later", 900), ]: with self.subTest(message=message): self.assertEqual(parse_upload_limit(limited(message), now).blocked_until, now + timedelta(seconds=seconds)) def test_structured_reset_and_retry_after_take_precedence(self): now = datetime(2026, 9, 17, tzinfo=timezone.utc) error = limited(retry_after=120) self.assertEqual(parse_upload_limit(error, now).blocked_until, now + timedelta(seconds=120)) error.body["detail"]["reset_at"] = "2026-09-18T00:00:00Z" self.assertEqual(parse_upload_limit(error, now).blocked_until, now + timedelta(days=1)) def test_other_failures_are_not_upload_quota(self): for error in [ UpstreamHTTPError("/backend-api/conversation", 429, {"detail": {"code": "throttled"}}), UpstreamHTTPError("/backend-api/files", 403, {"detail": {"code": "throttled"}}), UpstreamHTTPError("/backend-api/files", 429, {"detail": {"code": "server_busy"}}), RuntimeError("/backend-api/files status=429 throttled"), ]: self.assertIsNone(parse_upload_limit(error)) def test_json_string_body(self): import json error = limited() error.body = json.dumps(error.body) self.assertIsNotNone(parse_upload_limit(error)) def test_retry_after_http_date(self): from email.utils import format_datetime future = datetime.now(timezone.utc) + timedelta(hours=1) response = SimpleNamespace(status_code=429, text="", json=lambda: {}, headers={"Retry-After": format_datetime(future, usegmt=True)}) with self.assertRaises(UpstreamHTTPError) as caught: ensure_ok(response, "/backend-api/files") self.assertTrue(3590 <= caught.exception.retry_after <= 3600) class UploadPoolTests(unittest.TestCase): def setUp(self): self.temp = tempfile.TemporaryDirectory() self.addCleanup(self.temp.cleanup) for target in ("services.account_service.AccountService._save_cumulative_total", "services.account_service.log_service.add"): patcher = patch(target) patcher.start() self.addCleanup(patcher.stop) self.storage = JSONStorageBackend(Path(self.temp.name) / "accounts.json") self.service = AccountService(self.storage) self.service.add_account_items([ {"access_token": token, "status": "正常", "quota": 3, "type": "free"} for token in ("first", "second") ]) self.service.fetch_remote_info = lambda token, event: self.service.get_account(token) self.future = datetime.now(timezone.utc) + timedelta(hours=20) def block(self, token="first", until=None): self.service.mark_upload_limited(token, UploadLimit(until or self.future, "文件上传额度已用尽")) def test_cooling_survives_restart_and_refresh_without_changing_image_quota(self): self.block() self.service.update_account("first", {"status": "正常", "quota": 4}, quiet=True) restored = AccountService(self.storage) account = restored.get_account("first") self.assertEqual(account["upload_blocked_until"], self.future.isoformat()) self.assertEqual(account["quota"], 4) self.assertEqual(account["status"], "正常") self.assertNotIn("first", restored._list_ready_candidate_tokens(requires_upload=True)) self.assertIn("first", restored._list_ready_candidate_tokens()) def test_expired_limit_allows_upload_and_success_clears_it(self): self.block(until=datetime.now(timezone.utc) - timedelta(seconds=1)) token = self.service.get_available_access_token(requires_upload=True) self.assertEqual(token, "first") self.service.clear_expired_upload_limit(token) self.service.release_image_slot(token) self.assertNotIn("upload_blocked_until", self.service.get_account(token)) def test_success_cannot_clear_newer_concurrent_limit(self): self.block() self.service.clear_expired_upload_limit("first") self.assertEqual(self.service.get_account("first")["upload_blocked_until"], self.future.isoformat()) def test_all_blocked_raises_without_remote_requests_or_slot_leaks(self): self.block("first") self.block("second") with patch.object(self.service, "fetch_remote_info") as fetch: with self.assertRaises(UploadUnavailableError) as caught: self.service.get_available_access_token(requires_upload=True) fetch.assert_not_called() self.assertIn("最早预计恢复时间", str(caught.exception)) self.assertEqual(self.service._image_inflight, {}) def test_selection_excludes_rotated_token_alias(self): self.service._token_aliases["old-first"] = "first" token = self.service.get_available_access_token(requires_upload=True, excluded_tokens={"old-first"}) self.assertEqual(token, "second") self.service.release_image_slot(token) def test_limit_discovered_during_refresh_is_not_lost_to_stale_snapshot(self): def refresh(token, event): snapshot = self.service.get_account(token) if token == "first": self.block("first") return snapshot self.service.fetch_remote_info = refresh token = self.service.get_available_access_token(requires_upload=True) self.assertEqual(token, "second") self.service.release_image_slot(token) self.assertEqual(self.service._image_inflight, {}) def test_api_returns_upload_specific_429(self): from fastapi import FastAPI from fastapi.testclient import TestClient from api.ai import create_router self.block("first") self.block("second") app = FastAPI() app.include_router(create_router()) with patch.object(conversation, "account_service", self.service): with TestClient(app) as client: response = client.post( "/v1/images/edits", headers={"Authorization": f"Bearer {config.auth_key}"}, data={"model": "gpt-image-2", "prompt": "edit"}, files={"image": ("reference.png", b"reference", "image/png")}, ) self.assertEqual(response.status_code, 429, response.text) self.assertEqual(response.json()["error"]["code"], "upload_limit_exceeded") self.assertIn("预计恢复", response.json()["error"]["message"]) def test_busy_pool_obeys_deadline(self): import time with self.assertRaisesRegex(RuntimeError, "超时"): self.service.get_available_access_token(requires_upload=True, deadline=time.monotonic() - 1) self.assertEqual(self.service._image_inflight, {}) def run_request(self, stream, images=None, n=1): request = conversation.ConversationRequest(model="gpt-image-2", prompt="edit", images=images, n=n, progress_callback=MagicMock()) def backend(token=None, **kwargs): return SimpleNamespace(access_token=kwargs["access_token"], close=MagicMock()) with ( patch.object(conversation, "account_service", self.service), patch.object(conversation, "OpenAIBackendAPI", side_effect=backend), patch.object(conversation, "stream_image_outputs", side_effect=stream), ): outputs = list(conversation.stream_image_outputs_with_pool(request)) return outputs, request def test_first_account_throttled_second_succeeds_and_inputs_preserved(self): attempted = [] def stream(backend, request, index, total): attempted.append(backend.access_token) self.assertEqual(request.images, ["reference"]) self.assertEqual(request.prompt, "edit") if backend.access_token == "first": raise limited() return iter([conversation.ImageOutput(kind="result", model=request.model, index=index, total=total)]) result, request = self.run_request(stream, ["reference"]) self.assertEqual(attempted, ["first", "second"]) self.assertEqual(len(result), 1) self.assertEqual(self.service.get_account("first")["quota"], 3) self.assertEqual(self.service.get_account("second")["quota"], 2) self.assertIn("upload_blocked_until", self.service.get_account("first")) self.assertEqual(self.service._image_inflight, {}) request.progress_callback.assert_any_call("switching_upload_account") def test_all_accounts_fail_once_and_return_friendly_429(self): calls = [] def stream(backend, *args): calls.append(backend.access_token) raise limited() with self.assertRaises(conversation.ImageGenerationError) as caught: self.run_request(stream, ["reference"]) self.assertEqual(set(calls), {"first", "second"}) self.assertEqual(len(calls), 2) self.assertEqual(caught.exception.status_code, 429) self.assertEqual(caught.exception.code, "upload_limit_exceeded") self.assertNotIn("backend-api", str(caught.exception)) self.assertEqual(self.service._image_inflight, {}) def test_retry_cap_leaves_untried_accounts_for_later(self): self.service.add_account_items([ {"access_token": f"extra-{i}", "status": "正常", "quota": 3} for i in range(5) ]) calls = [] def stream(backend, *args): calls.append(backend.access_token) raise limited() with self.assertRaises(conversation.ImageGenerationError) as caught: self.run_request(stream, ["reference"]) self.assertEqual(len(set(calls)), 5) self.assertIn("自动切换已达到上限", str(caught.exception)) self.assertEqual(self.service._image_inflight, {}) def test_text_only_generation_still_uses_upload_limited_accounts(self): self.block("first") self.block("second") def stream(backend, request, index, total): return iter([conversation.ImageOutput(kind="result", model=request.model, index=index, total=total)]) result, _ = self.run_request(stream) self.assertEqual(len(result), 1) self.assertEqual(self.service._image_inflight, {}) def test_unrelated_429_does_not_switch_accounts(self): calls = [] def stream(backend, *args): calls.append(backend.access_token) raise UpstreamHTTPError("/backend-api/conversation", 429, {"detail": {"code": "throttled"}}) with self.assertRaises(conversation.ImageGenerationError): self.run_request(stream, ["reference"]) self.assertEqual(len(calls), 1) self.assertNotIn("upload_blocked_until", self.service.get_account(calls[0])) self.assertEqual(self.service._image_inflight, {}) def test_parallel_all_blocked_preserves_rate_limit_error(self): self.block("first") self.block("second") with patch.dict(config.data, {"image_parallel_generation": True}): with self.assertRaises(conversation.ImageGenerationError) as caught: self.run_request(MagicMock(), ["reference"], n=2) self.assertEqual(caught.exception.status_code, 429) self.assertEqual(self.service._image_inflight, {}) def test_successful_reference_upload_clears_only_after_all_files_uploaded(self): self.block(until=datetime.now(timezone.utc) - timedelta(seconds=1)) backend = object.__new__(OpenAIBackendAPI) backend.access_token = "first" backend._report_progress = MagicMock() backend._upload_image = MagicMock(return_value={"file_id": "test"}) backend._bootstrap = MagicMock(side_effect=RuntimeError("stop after upload")) with patch("services.openai_backend_api.account_service", self.service): with self.assertRaisesRegex(RuntimeError, "stop after upload"): list(backend._stream_picture_conversation("edit", "gpt-image-2", ["ref-1", "ref-2"])) self.assertEqual(backend._upload_image.call_count, 2) self.assertNotIn("upload_blocked_until", self.service.get_account("first")) if __name__ == "__main__": unittest.main()