119 lines
4.1 KiB
Python
119 lines
4.1 KiB
Python
import base64
|
|
import json
|
|
import unittest
|
|
from typing import Any
|
|
|
|
from services.account_service import AccountService
|
|
|
|
|
|
class MemoryStorage:
|
|
def __init__(self, accounts: list[dict[str, Any]] | None = None) -> None:
|
|
self.accounts = list(accounts or [])
|
|
|
|
def load_accounts(self) -> list[dict[str, Any]]:
|
|
return list(self.accounts)
|
|
|
|
def save_accounts(self, accounts: list[dict[str, Any]]) -> None:
|
|
self.accounts = list(accounts)
|
|
|
|
def load_auth_keys(self) -> list[dict[str, Any]]:
|
|
return []
|
|
|
|
def save_auth_keys(self, auth_keys: list[dict[str, Any]]) -> None:
|
|
pass
|
|
|
|
def health_check(self) -> dict[str, Any]:
|
|
return {"ok": True}
|
|
|
|
def get_backend_info(self) -> dict[str, Any]:
|
|
return {"type": "memory"}
|
|
|
|
|
|
def make_jwt(payload: dict[str, Any]) -> str:
|
|
def encode(value: dict[str, Any]) -> str:
|
|
raw = json.dumps(value, separators=(",", ":")).encode("utf-8")
|
|
return base64.urlsafe_b64encode(raw).decode("ascii").rstrip("=")
|
|
|
|
return f'{encode({"alg": "none", "typ": "JWT"})}.{encode(payload)}.sig'
|
|
|
|
|
|
class AccountExportTests(unittest.TestCase):
|
|
def test_build_export_items_uses_codex_shape_and_jwt_claims(self) -> None:
|
|
access_token = make_jwt(
|
|
{
|
|
"exp": 0,
|
|
"iat": 3600,
|
|
"https://api.openai.com/auth": {"chatgpt_account_id": "acct_123"},
|
|
"https://api.openai.com/profile": {"email": "test@example.com"},
|
|
}
|
|
)
|
|
id_token = make_jwt({"email": "fallback@example.com"})
|
|
service = AccountService(
|
|
MemoryStorage(
|
|
[
|
|
{
|
|
"access_token": access_token,
|
|
"id_token": id_token,
|
|
"refresh_token": "rt_test",
|
|
}
|
|
]
|
|
)
|
|
)
|
|
|
|
[item] = service.build_export_items([access_token])
|
|
|
|
self.assertEqual(item["type"], "codex")
|
|
self.assertEqual(item["email"], "test@example.com")
|
|
self.assertEqual(item["expired"], "1970-01-01T08:00:00+08:00")
|
|
self.assertEqual(item["account_id"], "acct_123")
|
|
self.assertEqual(item["access_token"], access_token)
|
|
self.assertEqual(item["last_refresh"], "1970-01-01T09:00:00+08:00")
|
|
self.assertEqual(item["id_token"], id_token)
|
|
self.assertEqual(item["refresh_token"], "rt_test")
|
|
|
|
def test_build_export_items_skips_accounts_missing_complete_tokens(self) -> None:
|
|
complete_access_token = make_jwt({"exp": 0})
|
|
complete_id_token = make_jwt({"email": "complete@example.com"})
|
|
service = AccountService(
|
|
MemoryStorage(
|
|
[
|
|
{"access_token": "only_access"},
|
|
{"access_token": "missing_id", "refresh_token": "rt_missing_id"},
|
|
{"access_token": complete_access_token, "id_token": complete_id_token, "refresh_token": "rt_complete"},
|
|
]
|
|
)
|
|
)
|
|
|
|
items = service.build_export_items()
|
|
|
|
self.assertEqual(len(items), 1)
|
|
self.assertEqual(items[0]["access_token"], complete_access_token)
|
|
self.assertEqual(items[0]["id_token"], complete_id_token)
|
|
self.assertEqual(items[0]["refresh_token"], "rt_complete")
|
|
|
|
def test_add_account_items_preserves_export_fields_without_overwriting_plan_type(self) -> None:
|
|
service = AccountService(MemoryStorage())
|
|
|
|
result = service.add_account_items(
|
|
[
|
|
{
|
|
"type": "codex",
|
|
"access_token": "access_token_test",
|
|
"refresh_token": "rt_test",
|
|
"account_id": "acct_123",
|
|
}
|
|
]
|
|
)
|
|
|
|
account = service.get_account("access_token_test")
|
|
self.assertEqual(result["added"], 1)
|
|
self.assertIsNotNone(account)
|
|
self.assertEqual(account["type"], "free")
|
|
self.assertEqual(account["export_type"], "codex")
|
|
self.assertEqual(account["refresh_token"], "rt_test")
|
|
self.assertEqual(account["account_id"], "acct_123")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|