from __future__ import annotations import os import tempfile import unittest from pathlib import Path from unittest.mock import patch os.environ.setdefault("CHATGPT2API_AUTH_KEY", "test-auth") from services.account_service import AccountService from services.auth_service import AuthService from services.config import config from services.openai_backend_api import InvalidAccessTokenError from services.storage.json_storage import JSONStorageBackend from utils.helper import anonymize_token, split_image_model class AccountCapabilityTests(unittest.TestCase): def test_image_accounts_require_positive_quota(self) -> None: self.assertFalse( AccountService._is_image_account_available( {"status": "限流", "quota": 1} ) ) self.assertFalse( AccountService._is_image_account_available( {"status": "正常", "quota": 0} ) ) self.assertTrue(AccountService._is_image_account_available({"status": "正常", "quota": 1})) def test_prolite_variants_are_normalized(self) -> None: with tempfile.TemporaryDirectory() as tmp_dir: service = AccountService(JSONStorageBackend(Path(tmp_dir) / "accounts.json")) self.assertEqual(service._normalize_account_type("prolite"), "ProLite") self.assertEqual(service._normalize_account_type("pro_lite"), "ProLite") def test_search_account_type_ignores_unrelated_scalar_values(self) -> None: with tempfile.TemporaryDirectory() as tmp_dir: service = AccountService(JSONStorageBackend(Path(tmp_dir) / "accounts.json")) self.assertIsNone( service._search_account_type( { "amr": ["pwd", "otp", "mfa"], "chatgpt_compute_residency": "no_constraint", "chatgpt_data_residency": "no_constraint", "user_id": "user-I52GFfLGFM0dokFk2dBiKEBn", } ) ) def test_mark_image_result_consumes_quota(self) -> None: with tempfile.TemporaryDirectory() as tmp_dir: service = AccountService(JSONStorageBackend(Path(tmp_dir) / "accounts.json")) service.add_accounts(["token-1"]) service.update_account( "token-1", { "status": "正常", "quota": 1, }, ) updated = service.mark_image_result("token-1", success=True) self.assertIsNotNone(updated) self.assertEqual(updated["quota"], 0) self.assertEqual(updated["status"], "限流") def test_split_image_model_supports_plan_type_prefix(self) -> None: self.assertEqual(split_image_model("gpt-image-2"), (None, "gpt-image-2")) self.assertEqual(split_image_model("plus-codex-gpt-image-2"), ("plus", "codex-gpt-image-2")) self.assertEqual(split_image_model("team-codex-gpt-image-2"), ("team", "codex-gpt-image-2")) self.assertEqual(split_image_model("pro-codex-gpt-image-2"), ("pro", "codex-gpt-image-2")) self.assertEqual(split_image_model("plus-gpt-image-2"), (None, None)) self.assertEqual(split_image_model("unknown-image-model"), (None, None)) def test_get_available_access_token_filters_by_plan_type(self) -> None: with tempfile.TemporaryDirectory() as tmp_dir: service = AccountService(JSONStorageBackend(Path(tmp_dir) / "accounts.json")) service.add_account_items( [ {"access_token": "token-plus", "type": "Plus", "status": "正常", "quota": 3}, {"access_token": "token-pro", "type": "Pro", "status": "正常", "quota": 3}, ] ) service.fetch_remote_info = lambda access_token, event="fetch_remote_info": service.get_account(access_token) plus_token = service.get_available_access_token(plan_type="plus") pro_token = service.get_available_access_token(plan_type="pro") service.release_image_slot(plus_token) service.release_image_slot(pro_token) self.assertEqual(plus_token, "token-plus") self.assertEqual(pro_token, "token-pro") def test_refresh_accounts_can_remove_invalid_token_without_confirmation_delay(self) -> None: original_value = config.data.get("auto_remove_invalid_accounts") config.data["auto_remove_invalid_accounts"] = True try: with tempfile.TemporaryDirectory() as tmp_dir: service = AccountService(JSONStorageBackend(Path(tmp_dir) / "accounts.json")) service.add_account_items([{"access_token": "invalid-token", "status": "正常"}]) with patch( "services.openai_backend_api.OpenAIBackendAPI.get_user_info", side_effect=InvalidAccessTokenError("token invalidated (/backend-api/me)"), ): result = service.refresh_accounts(["invalid-token"], defer_invalid_removal=False) self.assertEqual(result["refreshed"], 0) self.assertEqual(len(result["errors"]), 1) self.assertEqual(result["items"], []) self.assertIsNone(service.get_account("invalid-token")) finally: if original_value is None: config.data.pop("auto_remove_invalid_accounts", None) else: config.data["auto_remove_invalid_accounts"] = original_value def test_refresh_accounts_defers_invalid_token_removal_by_default(self) -> None: original_value = config.data.get("auto_remove_invalid_accounts") config.data["auto_remove_invalid_accounts"] = True try: with tempfile.TemporaryDirectory() as tmp_dir: service = AccountService(JSONStorageBackend(Path(tmp_dir) / "accounts.json")) service.add_account_items([{"access_token": "invalid-token", "status": "正常"}]) with patch( "services.openai_backend_api.OpenAIBackendAPI.get_user_info", side_effect=InvalidAccessTokenError("token invalidated (/backend-api/me)"), ): result = service.refresh_accounts(["invalid-token"]) account = service.get_account("invalid-token") self.assertEqual(result["refreshed"], 0) self.assertEqual(len(result["errors"]), 1) self.assertIsNotNone(account) self.assertEqual(account["invalid_count"], 1) finally: if original_value is None: config.data.pop("auto_remove_invalid_accounts", None) else: config.data["auto_remove_invalid_accounts"] = original_value class TokenLogTests(unittest.TestCase): def test_anonymize_token_hides_raw_value(self) -> None: token = "super-secret-token" token_ref = anonymize_token(token) self.assertTrue(token_ref.startswith("token:")) self.assertNotIn(token, token_ref) class AuthServiceTests(unittest.TestCase): def test_create_authenticate_disable_and_delete_user_key(self) -> None: with tempfile.TemporaryDirectory() as tmp_dir: service = AuthService(JSONStorageBackend(Path(tmp_dir) / "accounts.json", Path(tmp_dir) / "auth_keys.json")) item, raw_key = service.create_key(role="user", name="Alice") self.assertEqual(item["role"], "user") self.assertEqual(item["name"], "Alice") self.assertTrue(item["enabled"]) self.assertTrue(raw_key.startswith("sk-")) authed = service.authenticate(raw_key) self.assertIsNotNone(authed) self.assertEqual(authed["id"], item["id"]) self.assertEqual(authed["role"], "user") self.assertIsNotNone(authed["last_used_at"]) updated = service.update_key(item["id"], {"enabled": False}, role="user") self.assertIsNotNone(updated) self.assertFalse(updated["enabled"]) self.assertIsNone(service.authenticate(raw_key)) self.assertTrue(service.delete_key(item["id"], role="user")) self.assertFalse(service.delete_key(item["id"], role="user")) self.assertEqual(service.list_keys(role="user"), []) def test_authenticate_ignores_last_used_save_failure(self) -> None: with tempfile.TemporaryDirectory() as tmp_dir: service = AuthService(JSONStorageBackend(Path(tmp_dir) / "accounts.json", Path(tmp_dir) / "auth_keys.json")) item, raw_key = service.create_key(role="user", name="Alice") def fail_save() -> None: raise OSError("disk unavailable") service._save = fail_save authed = service.authenticate(raw_key) self.assertIsNotNone(authed) self.assertEqual(authed["id"], item["id"]) self.assertIsNotNone(authed["last_used_at"]) def test_update_user_key_replaces_raw_key(self) -> None: with tempfile.TemporaryDirectory() as tmp_dir: service = AuthService(JSONStorageBackend(Path(tmp_dir) / "accounts.json", Path(tmp_dir) / "auth_keys.json")) item, raw_key = service.create_key(role="user", name="Alice") updated = service.update_key(item["id"], {"key": "sk-user-custom-key"}, role="user") self.assertIsNotNone(updated) self.assertIsNone(service.authenticate(raw_key)) authed = service.authenticate("sk-user-custom-key") self.assertIsNotNone(authed) self.assertEqual(authed["id"], item["id"]) def test_user_key_name_must_be_unique(self) -> None: with tempfile.TemporaryDirectory() as tmp_dir: service = AuthService(JSONStorageBackend(Path(tmp_dir) / "accounts.json", Path(tmp_dir) / "auth_keys.json")) first, _ = service.create_key(role="user", name="Alice") second, _ = service.create_key(role="user", name="Bob") with self.assertRaisesRegex(ValueError, "这个名称已经在使用中了"): service.create_key(role="user", name="Alice") with self.assertRaisesRegex(ValueError, "这个名称已经在使用中了"): service.update_key(second["id"], {"name": "Alice"}, role="user") updated = service.update_key(first["id"], {"name": "Alice"}, role="user") self.assertIsNotNone(updated) self.assertEqual(updated["name"], "Alice") if __name__ == "__main__": unittest.main()