from __future__ import annotations import os import tempfile import unittest from pathlib import Path os.environ.setdefault("CHATGPT2API_AUTH_KEY", "test-auth") from services.account_service import AccountService from services.auth_service import AuthService from services.storage.json_storage import JSONStorageBackend from utils.helper import anonymize_token, split_image_model class AccountCapabilityTests(unittest.TestCase): def test_unknown_quota_accounts_are_available_only_when_not_throttled(self) -> None: self.assertFalse( AccountService._is_image_account_available( {"status": "限流", "image_quota_unknown": True, "quota": 0} ) ) self.assertTrue( AccountService._is_image_account_available( {"status": "正常", "image_quota_unknown": True, "quota": 0} ) ) 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_does_not_consume_unknown_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": 0, "image_quota_unknown": True, }, ) updated = service.mark_image_result("token-1", success=True) self.assertIsNotNone(updated) self.assertEqual(updated["quota"], 0) self.assertEqual(updated["status"], "正常") self.assertTrue(updated["image_quota_unknown"]) 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") 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()