from __future__ import annotations import tempfile import unittest from pathlib import Path from unittest.mock import patch from fastapi import FastAPI from fastapi.testclient import TestClient import api.accounts as accounts_api from services.account_service import AccountService from services.config import config from services.storage.json_storage import JSONStorageBackend class ManualAccountStatusTests(unittest.TestCase): def setUp(self): temp = tempfile.TemporaryDirectory() self.addCleanup(temp.cleanup) self.storage = JSONStorageBackend(Path(temp.name) / "accounts.json") self.storage.save_accounts([ {"access_token": "token", "status": "正常", "quota": 3}, ]) self.service = AccountService(self.storage) settings = patch.dict(config.data, { "auto_remove_rate_limited_accounts": False, "auto_remove_invalid_accounts": False, "auto_relogin_after_refresh": False, }) settings.start() self.addCleanup(settings.stop) logs = patch("services.account_service.log_service.add") logs.start() self.addCleanup(logs.stop) def refresh(self, during_fetch=None): def remote_info(_backend): if during_fetch: during_fetch() return {"status": "正常", "quota": 5, "email": "test@example.test"} with patch.object(self.service, "refresh_access_token", side_effect=lambda token, **_: token), patch( "services.openai_backend_api.OpenAIBackendAPI.get_user_info", autospec=True, side_effect=remote_info ): return self.service.fetch_remote_info("token", "test") def test_refresh_finishing_after_manual_edit_preserves_status(self): for status in ("禁用", "限流"): with self.subTest(status=status): # Model the admin edit happening while the remote request is running. result = self.refresh(lambda: self.service.update_account("token", {"status": status}, manual=True)) self.assertEqual(result["status"], status) self.assertEqual(result["quota"], 5) self.assertEqual(result["email"], "test@example.test") self.assertFalse(self.service._is_image_account_available(result)) def test_task_completion_does_not_clear_manual_status_or_delete_account(self): with patch.dict(config.data, {"auto_remove_rate_limited_accounts": True}): for status in ("禁用", "限流"): for quota in (1, 3): with self.subTest(status=status, quota=quota): self.service.update_account("token", {"status": status, "quota": quota}, manual=True) result = self.service.mark_image_result("token", True) self.assertIsNotNone(result) self.assertEqual(result["status"], status) self.assertEqual(result["quota"], quota - 1) self.assertIsNotNone(self.service.get_account("token")) def test_manual_status_survives_reload_token_rotation_and_background_update(self): for status in ("禁用", "限流"): with self.subTest(status=status): self.service.update_account("token", {"status": status}, manual=True) reloaded = AccountService(self.storage) new_token = reloaded._apply_refreshed_tokens("token", {"access_token": "rotated"}, "test") result = reloaded.update_account(new_token, {"status": "正常"}, quiet=True) self.assertEqual(result["status"], status) reloaded._apply_refreshed_tokens(new_token, {"access_token": "token"}, "test") def test_manual_normal_clears_override_and_allows_automatic_recovery(self): self.service.update_account("token", {"status": "限流"}, manual=True) self.assertNotIn("token", self.service.list_limited_tokens()) self.service.update_account("token", {"status": "正常"}, manual=True) self.service.update_account("token", {"status": "限流", "quota": 0}) self.assertIn("token", self.service.list_limited_tokens()) self.assertEqual(self.refresh()["status"], "正常") self.assertIsNone(self.service.get_account("token")["manual_status"]) def test_legacy_disabled_accounts_stay_disabled_but_automatic_limits_recover(self): for status in ("禁用", "限流"): with self.subTest(status=status): self.storage.save_accounts([{"access_token": "token", "status": status, "quota": 0}]) self.service = AccountService(self.storage) expected = "禁用" if status == "禁用" else "正常" self.assertEqual(self.refresh()["status"], expected) def test_background_invalid_token_cleanup_keeps_manual_accounts(self): with patch.dict(config.data, {"auto_remove_invalid_accounts": True}): for status in ("禁用", "限流"): with self.subTest(status=status): self.service.update_account("token", {"status": status}, manual=True) self.assertFalse(self.service.remove_invalid_token("token", "test")) self.assertEqual(self.service.get_account("token")["status"], status) def test_admin_update_endpoint_sets_and_clears_manual_override(self): app = FastAPI() app.include_router(accounts_api.create_router()) with patch.object(accounts_api, "account_service", self.service), patch.object( accounts_api, "require_admin", return_value={"id": "admin", "role": "admin"} ), TestClient(app) as client: for status in ("禁用", "限流", "正常"): response = client.post("/api/accounts/update", json={"access_token": "token", "status": status}) self.assertEqual(response.status_code, 200, response.text) item = response.json()["item"] self.assertEqual(item["manual_status"], None if status == "正常" else status) client.post("/api/accounts/update", json={"access_token": "token", "status": "禁用"}) response = client.post("/api/accounts/update", json={"access_token": "token", "quota": 8}) self.assertEqual(response.json()["item"]["manual_status"], "禁用") def test_manual_edit_during_invalid_cleanup_prevents_deletion(self): delete = self.service.delete_accounts def concurrent_delete(tokens, **kwargs): self.service.update_account("token", {"status": "禁用"}, manual=True) return delete(tokens, **kwargs) with patch.dict(config.data, {"auto_remove_invalid_accounts": True}), patch.object( self.service, "delete_accounts", side_effect=concurrent_delete ): self.assertFalse(self.service.remove_invalid_token("token", "test")) self.assertEqual(self.service.get_account("token")["status"], "禁用") # Explicit administrator deletion is still allowed. self.assertEqual(self.service.delete_accounts(["token"])["removed"], 1) if __name__ == "__main__": unittest.main()