Initial project import

test/test_account_live_validation.py

@@ -0,0 +1,94 @@
+import tempfile
+import unittest
+from datetime import datetime, timedelta, timezone
+from pathlib import Path
+from unittest.mock import patch
+
+from services.account_service import AccountService
+from services.storage.json_storage import JSONStorageBackend
+from services.upload_limits import UploadLimit, UploadUnavailableError
+
+
+class AccountLiveValidationTests(unittest.TestCase):
+    def setUp(self):
+        temp = tempfile.TemporaryDirectory()
+        self.addCleanup(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.service = AccountService(JSONStorageBackend(Path(temp.name) / "accounts.json"))
+        self.service.add_account_items([{"access_token": "test-token", "quota": 5, "status": "正常"}])
+        patcher = patch("services.openai_backend_api.OpenAIBackendAPI")
+        self.backend = patcher.start()
+        self.addCleanup(patcher.stop)
+        self.backend.return_value.get_user_info.return_value = {"quota": 5, "status": "正常"}
+
+    def acquire(self, **kwargs):
+        token = self.service.get_available_access_token(**kwargs)
+        self.service.release_image_slot(token)
+        return token
+
+    def test_repeated_selection_always_refreshes_remote_quota(self):
+        self.acquire()
+        self.service.mark_image_result("test-token", True)
+        self.backend.return_value.get_user_info.return_value = {"quota": 2, "status": "正常"}
+        self.acquire()
+        self.assertEqual(self.backend.return_value.get_user_info.call_count, 2)
+        self.assertEqual(self.service.get_account("test-token")["quota"], 2)
+
+    def test_recent_background_refresh_does_not_skip_live_validation(self):
+        self.service.fetch_remote_info("test-token")
+        self.acquire()
+        self.assertEqual(self.backend.return_value.get_user_info.call_count, 2)
+
+    def test_remote_quota_exhaustion_is_detected_immediately(self):
+        self.acquire()
+        self.backend.return_value.get_user_info.return_value = {"quota": 0, "status": "限流"}
+        with self.assertRaises(RuntimeError):
+            self.acquire()
+        self.assertEqual(self.backend.return_value.get_user_info.call_count, 2)
+        self.assertEqual(self.service._image_inflight, {})
+
+    def test_generation_failure_is_followed_by_live_validation(self):
+        self.acquire()
+        self.service.mark_image_result("test-token", False)
+        self.acquire()
+        self.assertEqual(self.backend.return_value.get_user_info.call_count, 2)
+
+    def test_manual_disable_and_zero_quota_are_excluded_before_validation(self):
+        self.acquire()
+        self.service.update_account("test-token", {"status": "禁用"}, manual=True)
+        with self.assertRaises(RuntimeError):
+            self.acquire()
+        self.service.update_account("test-token", {"status": "正常", "quota": 0}, manual=True)
+        with self.assertRaises(RuntimeError):
+            self.acquire()
+        self.assertEqual(self.backend.return_value.get_user_info.call_count, 1)
+
+    def test_upload_cooldown_is_respected_but_text_generation_revalidates(self):
+        self.acquire()
+        self.service.mark_upload_limited("test-token", UploadLimit(datetime.now(timezone.utc) + timedelta(hours=1), "limited"))
+        with self.assertRaises(UploadUnavailableError):
+            self.acquire(requires_upload=True)
+        self.assertEqual(self.acquire(), "test-token")
+        self.assertEqual(self.backend.return_value.get_user_info.call_count, 2)
+
+    def test_remote_failure_does_not_fall_back_to_previous_validation(self):
+        self.acquire()
+        self.backend.return_value.get_user_info.side_effect = RuntimeError("upstream unavailable")
+        with self.assertRaises(RuntimeError):
+            self.acquire()
+        self.assertEqual(self.backend.return_value.get_user_info.call_count, 2)
+        self.assertEqual(self.service._image_inflight, {})
+
+    def test_stats_count_active_upload_limits_independently(self):
+        self.service.add_account_items([{"access_token": "expired", "quota": 5, "status": "正常"}])
+        now = datetime.now(timezone.utc)
+        self.service.mark_upload_limited("test-token", UploadLimit(now + timedelta(hours=1), "limited"))
+        self.service.mark_upload_limited("expired", UploadLimit(now - timedelta(seconds=1), "expired"))
+        stats = self.service.get_stats()
+        self.assertEqual(stats["upload_limited"], 1)
+        self.assertEqual(stats["active"], 2)
+        self.assertEqual(stats["limited"], 0)