234 lines
10 KiB
Python
234 lines
10 KiB
Python
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()
|