Initial project import

2026-09-18 03:46:58 +08:00
commit b600da4b76
209 changed files with 45904 additions and 0 deletions
@@ -0,0 +1,274 @@
+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()