95 lines
4.6 KiB
Python
95 lines
4.6 KiB
Python
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)
|