138 lines
7.1 KiB
Python
138 lines
7.1 KiB
Python
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()
|