| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227 |
- import sys
- import types
- import unittest
- from unittest.mock import patch
- import chatgpt_flow
- from config import AppConfig
- class _FakeGeo:
- locale = "en-US"
- timezone_id = "America/New_York"
- class _FakePage:
- def on(self, *_args, **_kwargs):
- pass
- class _FakeBrowserContext:
- def new_page(self):
- return _FakePage()
- class _FakeBrowser:
- def new_context(self, **_kwargs):
- return _FakeBrowserContext()
- def close(self):
- pass
- class _FakeChromium:
- def launch(self, **_kwargs):
- return _FakeBrowser()
- class _FakePlaywright:
- chromium = _FakeChromium()
- class _FakeSyncPlaywright:
- def __enter__(self):
- return _FakePlaywright()
- def __exit__(self, *_args):
- return False
- def _fake_sync_playwright():
- return _FakeSyncPlaywright()
- class FullRunAttemptLabelTests(unittest.TestCase):
- def test_retry_keeps_account_number_based_on_success_count(self):
- calls = []
- def fake_run_one_account(_full_ctx, _page, idx, total):
- calls.append((idx, total))
- if len(calls) == 1:
- raise RuntimeError("signup failed")
- return {"stage": "cpa_uploaded", "email": "ok@example.com"}
- fake_patchright = types.ModuleType("patchright")
- fake_sync_api = types.ModuleType("patchright.sync_api")
- fake_sync_api.sync_playwright = _fake_sync_playwright
- fake_patchright.sync_api = fake_sync_api
- logs = []
- cfg = AppConfig(account_count=1, cpa_url="https://cpa.example", cpa_management_key="token", headless=True)
- with patch.dict(sys.modules, {"patchright": fake_patchright, "patchright.sync_api": fake_sync_api}):
- with patch("geo_fingerprint.detect_openai_geo_fingerprint", return_value=_FakeGeo()):
- with patch.object(chatgpt_flow, "_run_one_account", side_effect=fake_run_one_account):
- chatgpt_flow.run_full(cfg, log=logs.append)
- self.assertEqual(calls, [(1, 1), (1, 1)])
- self.assertTrue(any("启动账号 1/1 第2次尝试" in line for line in logs))
- class SsoCpaOauthRouteTests(unittest.TestCase):
- def test_full_sso_account_uses_cpa_oauth_instead_of_direct_session_upload(self):
- cfg = AppConfig(
- signup_mode="sso",
- cpa_url="https://cpa.example/management.html#/oauth",
- cpa_management_key="secret",
- )
- ctx = chatgpt_flow.FullRunContext(cfg=cfg, log=lambda *_: None)
- signup_session = {
- "accessToken": "tok",
- "account": {"planType": "free"},
- }
- oauth_result = {
- "fileName": "codex-oauth.json",
- "email": "n@example.com",
- "planType": "codex-oauth",
- "status": "Authentication successful!",
- }
- with patch.object(chatgpt_flow, "signup_chatgpt_sso", return_value={
- "email": "n@example.com",
- "password": "pw",
- "session": signup_session,
- }):
- with patch.object(chatgpt_flow, "upload_session_to_cpa") as direct_upload:
- with patch.object(chatgpt_flow, "authorize_codex_oauth_to_cpa", return_value=oauth_result, create=True) as oauth_upload:
- with patch.object(chatgpt_flow, "upsert_account"):
- with patch.object(chatgpt_flow, "add_event"):
- record = chatgpt_flow._run_one_account(ctx, _FakePage(), 1, 1)
- self.assertEqual(record["stage"], "cpa_uploaded")
- direct_upload.assert_not_called()
- oauth_upload.assert_called_once()
- def test_full_sso_account_does_not_require_session_access_token_before_cpa_oauth(self):
- cfg = AppConfig(
- signup_mode="sso",
- cpa_url="https://cpa.example/management.html#/oauth",
- cpa_management_key="secret",
- )
- ctx = chatgpt_flow.FullRunContext(cfg=cfg, log=lambda *_: None)
- oauth_result = {
- "fileName": "codex-oauth.json",
- "email": "n@example.com",
- "planType": "codex-oauth",
- "status": "Authentication successful!",
- }
- with patch.object(chatgpt_flow, "signup_chatgpt_sso", return_value={
- "email": "n@example.com",
- "password": "pw",
- "session": {},
- }):
- with patch.object(chatgpt_flow, "upload_session_to_cpa") as direct_upload:
- with patch.object(chatgpt_flow, "authorize_codex_oauth_to_cpa", return_value=oauth_result, create=True) as oauth_upload:
- with patch.object(chatgpt_flow, "upsert_account"):
- with patch.object(chatgpt_flow, "add_event"):
- record = chatgpt_flow._run_one_account(ctx, _FakePage(), 1, 1)
- self.assertEqual(record["stage"], "cpa_uploaded")
- self.assertEqual(record["planType"], "codex-oauth")
- direct_upload.assert_not_called()
- oauth_upload.assert_called_once()
- self.assertEqual(oauth_upload.call_args.kwargs["email_hint"], "n@example.com")
- def test_sso_batch_uses_cpa_oauth_instead_of_direct_session_upload(self):
- fake_patchright = types.ModuleType("patchright")
- fake_sync_api = types.ModuleType("patchright.sync_api")
- fake_sync_api.sync_playwright = _fake_sync_playwright
- fake_patchright.sync_api = fake_sync_api
- signup_session = {
- "accessToken": "tok",
- "account": {"planType": "free"},
- }
- oauth_result = {
- "fileName": "codex-oauth.json",
- "email": "n@example.com",
- "planType": "codex-oauth",
- "status": "Authentication successful!",
- }
- with patch.dict(sys.modules, {"patchright": fake_patchright, "patchright.sync_api": fake_sync_api}):
- with patch("geo_fingerprint.detect_openai_geo_fingerprint", return_value=_FakeGeo()):
- with patch.object(chatgpt_flow, "signup_chatgpt_sso", return_value={
- "email": "n@example.com",
- "password": "pw",
- "session": signup_session,
- }):
- with patch.object(chatgpt_flow, "upload_session_to_cpa") as direct_upload:
- with patch.object(chatgpt_flow, "authorize_codex_oauth_to_cpa", return_value=oauth_result, create=True) as oauth_upload:
- with patch.object(chatgpt_flow, "upsert_account"):
- with patch.object(chatgpt_flow, "add_event"):
- ctx = chatgpt_flow.run_sso_batch(
- account_count=1,
- cpa_url="https://cpa.example/management.html#/oauth",
- cpa_management_key="secret",
- headless=True,
- log=lambda *_: None,
- )
- self.assertEqual(ctx.accounts[0]["stage"], "cpa_uploaded")
- direct_upload.assert_not_called()
- oauth_upload.assert_called_once()
- self.assertEqual(oauth_upload.call_args.kwargs["cpa_url"], "https://cpa.example/management.html#/oauth")
- self.assertEqual(oauth_upload.call_args.kwargs["management_key"], "secret")
- def test_sso_batch_does_not_require_session_access_token_before_cpa_oauth(self):
- fake_patchright = types.ModuleType("patchright")
- fake_sync_api = types.ModuleType("patchright.sync_api")
- fake_sync_api.sync_playwright = _fake_sync_playwright
- fake_patchright.sync_api = fake_sync_api
- oauth_result = {
- "fileName": "codex-oauth.json",
- "email": "n@example.com",
- "planType": "codex-oauth",
- "status": "Authentication successful!",
- }
- with patch.dict(sys.modules, {"patchright": fake_patchright, "patchright.sync_api": fake_sync_api}):
- with patch("geo_fingerprint.detect_openai_geo_fingerprint", return_value=_FakeGeo()):
- with patch.object(chatgpt_flow, "signup_chatgpt_sso", return_value={
- "email": "n@example.com",
- "password": "pw",
- "session": {},
- }):
- with patch.object(chatgpt_flow, "upload_session_to_cpa") as direct_upload:
- with patch.object(chatgpt_flow, "authorize_codex_oauth_to_cpa", return_value=oauth_result, create=True) as oauth_upload:
- with patch.object(chatgpt_flow, "upsert_account"):
- with patch.object(chatgpt_flow, "add_event"):
- ctx = chatgpt_flow.run_sso_batch(
- account_count=1,
- cpa_url="https://cpa.example/management.html#/oauth",
- cpa_management_key="secret",
- headless=True,
- log=lambda *_: None,
- )
- self.assertEqual(ctx.accounts[0]["stage"], "cpa_uploaded")
- self.assertEqual(ctx.accounts[0]["planType"], "codex-oauth")
- direct_upload.assert_not_called()
- oauth_upload.assert_called_once()
- if __name__ == "__main__":
- unittest.main()
|