Initial project import

test/test_account_export.py

@@ -0,0 +1,118 @@
+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()