test_cpa_oauth.py 22 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617
  1. import unittest
  2. from unittest.mock import patch
  3. class CallbackUrlTests(unittest.TestCase):
  4. def test_accepts_localhost_auth_and_codex_callbacks_with_code_and_state(self):
  5. import cpa_oauth
  6. self.assertTrue(cpa_oauth.is_localhost_oauth_callback_url("http://localhost:1455/auth/callback?code=abc&state=xyz"))
  7. self.assertTrue(cpa_oauth.is_localhost_oauth_callback_url("http://127.0.0.1:1455/codex/callback?code=abc&state=xyz"))
  8. def test_rejects_non_localhost_or_missing_oauth_params(self):
  9. import cpa_oauth
  10. self.assertFalse(cpa_oauth.is_localhost_oauth_callback_url("https://example.com/auth/callback?code=abc&state=xyz"))
  11. self.assertFalse(cpa_oauth.is_localhost_oauth_callback_url("http://localhost:1455/auth/callback?code=abc"))
  12. self.assertFalse(cpa_oauth.is_localhost_oauth_callback_url("http://localhost:1455/other?code=abc&state=xyz"))
  13. class FetchOauthUrlTests(unittest.TestCase):
  14. def test_waits_for_oauth_panel_after_management_login(self):
  15. import cpa_oauth
  16. class _FakeLocator:
  17. def __init__(self, page, selector):
  18. self.page = page
  19. self.selector = selector
  20. def count(self):
  21. if "authUrlValue" in self.selector:
  22. if self.page.state == "oauth_generating":
  23. self.page.auth_url_polls += 1
  24. if self.page.auth_url_polls >= 2:
  25. self.page.state = "auth_url"
  26. return 1 if self.page.state == "auth_url" else 0
  27. if "password" in self.selector:
  28. return 1 if self.page.state == "login" else 0
  29. if 'has-text("Login")' in self.selector or 'has-text("登录")' in self.selector:
  30. return 1 if self.page.state in ("login", "oauth_card") else 0
  31. return 0
  32. def nth(self, _idx):
  33. return self
  34. def is_visible(self):
  35. return True
  36. def is_enabled(self):
  37. return True
  38. def fill(self, value):
  39. self.page.filled.append((self.selector, value))
  40. def click(self, **_kwargs):
  41. self.page.clicked.append((self.page.state, self.selector))
  42. if self.page.state == "login":
  43. self.page.state = "loading_after_login"
  44. elif self.page.state == "oauth_card":
  45. self.page.state = "oauth_generating"
  46. def inner_text(self, **_kwargs):
  47. if self.page.state == "auth_url":
  48. return "https://auth.openai.com/oauth/authorize?client_id=codex"
  49. return ""
  50. class _FakePage:
  51. def __init__(self):
  52. self.url = ""
  53. self.state = "new"
  54. self.wait_calls = 0
  55. self.auth_url_polls = 0
  56. self.clicked = []
  57. self.filled = []
  58. def goto(self, url, **_kwargs):
  59. self.url = url.rsplit("#", 1)[0] + "#/login"
  60. self.state = "login"
  61. def wait_for_timeout(self, *_args, **_kwargs):
  62. self.wait_calls += 1
  63. if self.state == "loading_after_login" and self.wait_calls >= 3:
  64. self.url = self.url.rsplit("#", 1)[0] + "#/oauth"
  65. self.state = "oauth_card"
  66. def locator(self, selector):
  67. return _FakeLocator(self, selector)
  68. page = _FakePage()
  69. oauth_url = cpa_oauth.fetch_cpa_oauth_url(
  70. page,
  71. cpa_url="https://cpa.example/management.html#/oauth",
  72. management_key="secret",
  73. log=lambda *_: None,
  74. timeout_sec=2,
  75. )
  76. self.assertTrue(oauth_url.startswith("https://auth.openai.com/oauth/authorize"))
  77. self.assertTrue(any(state == "login" for state, _selector in page.clicked))
  78. self.assertTrue(any(state == "oauth_card" for state, _selector in page.clicked))
  79. def test_reopens_oauth_route_when_management_login_lands_on_home(self):
  80. import cpa_oauth
  81. class _FakeLocator:
  82. def __init__(self, page, selector):
  83. self.page = page
  84. self.selector = selector
  85. def count(self):
  86. if "authUrlValue" in self.selector:
  87. if self.page.state == "oauth_generating":
  88. self.page.state = "auth_url"
  89. return 1 if self.page.state == "auth_url" else 0
  90. if "password" in self.selector:
  91. return 1 if self.page.state == "login" else 0
  92. if 'has-text("Login")' in self.selector or 'has-text("登录")' in self.selector:
  93. return 1 if self.page.state in ("login", "oauth_card") else 0
  94. return 0
  95. def nth(self, _idx):
  96. return self
  97. def is_visible(self):
  98. return True
  99. def is_enabled(self):
  100. return True
  101. def fill(self, value):
  102. self.page.filled.append((self.selector, value))
  103. def click(self, **_kwargs):
  104. self.page.clicked.append((self.page.state, self.selector))
  105. if self.page.state == "login":
  106. self.page.url = "https://cpa.example/management.html#/"
  107. self.page.state = "home"
  108. elif self.page.state == "oauth_card":
  109. self.page.state = "oauth_generating"
  110. def inner_text(self, **_kwargs):
  111. if self.page.state == "auth_url":
  112. return "https://auth.openai.com/oauth/authorize?client_id=codex"
  113. return ""
  114. class _FakePage:
  115. def __init__(self):
  116. self.url = ""
  117. self.state = "new"
  118. self.goto_calls = []
  119. self.clicked = []
  120. self.filled = []
  121. def goto(self, url, **_kwargs):
  122. self.goto_calls.append(url)
  123. self.url = url
  124. if len(self.goto_calls) == 1:
  125. self.url = "https://cpa.example/management.html#/login"
  126. self.state = "login"
  127. elif url.endswith("#/oauth"):
  128. self.state = "oauth_card"
  129. def wait_for_timeout(self, *_args, **_kwargs):
  130. pass
  131. def locator(self, selector):
  132. return _FakeLocator(self, selector)
  133. page = _FakePage()
  134. oauth_url = cpa_oauth.fetch_cpa_oauth_url(
  135. page,
  136. cpa_url="https://cpa.example/management.html#/oauth",
  137. management_key="secret",
  138. log=lambda *_: None,
  139. timeout_sec=2,
  140. )
  141. self.assertTrue(oauth_url.startswith("https://auth.openai.com/oauth/authorize"))
  142. self.assertGreaterEqual(page.goto_calls.count("https://cpa.example/management.html#/oauth"), 2)
  143. self.assertTrue(any(state == "home" for state, _selector in page.clicked) is False)
  144. class SubmitCallbackTests(unittest.TestCase):
  145. def test_submit_callback_fills_url_and_clicks_submit(self):
  146. import cpa_oauth
  147. class _FakeLocator:
  148. def __init__(self, page, selector):
  149. self.page = page
  150. self.selector = selector
  151. def count(self):
  152. if self.selector == 'input[placeholder*="localhost"]':
  153. return 1
  154. if "Submit Callback URL" in self.selector:
  155. return 1
  156. return 0
  157. def nth(self, _idx):
  158. return self
  159. def is_visible(self):
  160. return True
  161. def is_enabled(self):
  162. return True
  163. @property
  164. def first(self):
  165. return self
  166. def click(self, **_kwargs):
  167. self.page.clicked.append(self.selector)
  168. def fill(self, value):
  169. self.page.filled.append((self.selector, value))
  170. class _FakePage:
  171. def __init__(self):
  172. self.url = ""
  173. self.clicked = []
  174. self.filled = []
  175. self.goto_calls = []
  176. def goto(self, url, **_kwargs):
  177. self.url = url
  178. self.goto_calls.append(url)
  179. def locator(self, selector):
  180. return _FakeLocator(self, selector)
  181. def wait_for_timeout(self, *_args, **_kwargs):
  182. pass
  183. page = _FakePage()
  184. callback_url = "http://localhost:1455/auth/callback?code=abc&state=xyz"
  185. cpa_oauth.submit_oauth_callback_to_cpa(
  186. page,
  187. cpa_url="https://cpa.example/management.html#/oauth",
  188. callback_url=callback_url,
  189. log=lambda *_: None,
  190. wait_for_success=False,
  191. )
  192. self.assertEqual(page.goto_calls, ["https://cpa.example/management.html#/oauth"])
  193. self.assertEqual(page.filled, [('input[placeholder*="localhost"]', callback_url)])
  194. self.assertTrue(any("Submit Callback URL" in selector for selector in page.clicked))
  195. class CaptureCallbackTests(unittest.TestCase):
  196. def test_fills_email_hint_when_openai_oauth_shows_email_prompt(self):
  197. import cpa_oauth
  198. callback_url = "http://localhost:1455/auth/callback?code=abc&state=xyz"
  199. class _FakeLocator:
  200. def __init__(self, page, selector):
  201. self.page = page
  202. self.selector = selector
  203. def count(self):
  204. if 'input[type="email"]' in self.selector:
  205. return 1 if self.page.state == "email_prompt" else 0
  206. if 'button[type="submit"]' in self.selector:
  207. return 1 if self.page.state == "email_prompt" else 0
  208. return 0
  209. def nth(self, _idx):
  210. return self
  211. def is_visible(self):
  212. return True
  213. def is_enabled(self):
  214. return True
  215. def fill(self, value):
  216. self.page.filled.append((self.selector, value))
  217. def click(self, **_kwargs):
  218. self.page.clicked.append(self.selector)
  219. self.page.url = callback_url
  220. self.page.state = "callback"
  221. class _FakePage:
  222. def __init__(self):
  223. self.url = ""
  224. self.state = "new"
  225. self.filled = []
  226. self.clicked = []
  227. def goto(self, url, **_kwargs):
  228. self.url = url
  229. self.state = "email_prompt"
  230. def wait_for_timeout(self, *_args, **_kwargs):
  231. pass
  232. def locator(self, selector):
  233. return _FakeLocator(self, selector)
  234. page = _FakePage()
  235. captured = cpa_oauth.approve_openai_oauth_and_capture_callback(
  236. page,
  237. "https://auth.openai.com/oauth/authorize?client_id=codex",
  238. email_hint="n@example.com",
  239. log=lambda *_: None,
  240. )
  241. self.assertEqual(captured, callback_url)
  242. self.assertEqual(page.filled, [('input[type="email"]', "n@example.com")])
  243. self.assertIn('button[type="submit"]', page.clicked)
  244. def test_clicks_delayed_codex_consent_continue_after_email_continue(self):
  245. import cpa_oauth
  246. callback_url = "http://localhost:1455/auth/callback?code=abc&state=xyz"
  247. class _FakeLocator:
  248. def __init__(self, page, selector):
  249. self.page = page
  250. self.selector = selector
  251. def count(self):
  252. if 'input[type="email"]' in self.selector:
  253. return 1 if self.page.state == "email_prompt" else 0
  254. if 'button[type="submit"]' in self.selector:
  255. if self.page.state == "email_prompt":
  256. return 1
  257. if self.page.state == "loading_consent":
  258. self.page.button_polls += 1
  259. if self.page.button_polls >= 2:
  260. self.page.state = "consent"
  261. self.page.url = "https://auth.openai.com/sign-in-with-chatgpt/codex/consent"
  262. return 1
  263. return 1 if self.page.state == "consent" else 0
  264. return 0
  265. def nth(self, _idx):
  266. return self
  267. def is_visible(self):
  268. return True
  269. def is_enabled(self):
  270. return True
  271. def fill(self, value):
  272. self.page.filled.append((self.selector, value))
  273. def click(self, **_kwargs):
  274. self.page.clicked.append((self.page.state, self.selector))
  275. if self.page.state == "email_prompt":
  276. self.page.url = "https://auth.openai.com/log-in"
  277. self.page.state = "loading_consent"
  278. elif self.page.state == "consent":
  279. self.page.url = callback_url
  280. self.page.state = "callback"
  281. class _FakePage:
  282. def __init__(self):
  283. self.url = ""
  284. self.state = "new"
  285. self.button_polls = 0
  286. self.filled = []
  287. self.clicked = []
  288. def goto(self, url, **_kwargs):
  289. self.url = url
  290. self.state = "email_prompt"
  291. def wait_for_timeout(self, *_args, **_kwargs):
  292. pass
  293. def locator(self, selector):
  294. return _FakeLocator(self, selector)
  295. page = _FakePage()
  296. captured = cpa_oauth.approve_openai_oauth_and_capture_callback(
  297. page,
  298. "https://auth.openai.com/oauth/authorize?client_id=codex",
  299. email_hint="n@example.com",
  300. log=lambda *_: None,
  301. timeout_sec=1,
  302. )
  303. self.assertEqual(captured, callback_url)
  304. self.assertIn(("email_prompt", 'button[type="submit"]'), page.clicked)
  305. self.assertIn(("consent", 'button[type="submit"]'), page.clicked)
  306. def test_submits_email_only_once_when_prompt_remains_visible_before_consent(self):
  307. import cpa_oauth
  308. callback_url = "http://localhost:1455/auth/callback?code=abc&state=xyz"
  309. class _FakeLocator:
  310. def __init__(self, page, selector):
  311. self.page = page
  312. self.selector = selector
  313. def count(self):
  314. if 'input[type="email"]' in self.selector:
  315. return 1 if self.page.state in ("email_prompt", "loading_consent") else 0
  316. if 'button[type="submit"]' in self.selector:
  317. return 1 if self.page.state in ("email_prompt", "loading_consent", "consent") else 0
  318. return 0
  319. def nth(self, _idx):
  320. return self
  321. def is_visible(self):
  322. return True
  323. def is_enabled(self):
  324. return True
  325. def fill(self, value):
  326. self.page.filled.append((self.page.state, self.selector, value))
  327. def click(self, **_kwargs):
  328. self.page.clicked.append((self.page.state, self.selector))
  329. if self.page.state == "email_prompt":
  330. self.page.state = "loading_consent"
  331. self.page.url = "https://auth.openai.com/log-in"
  332. elif self.page.state == "loading_consent":
  333. self.page.state = "consent"
  334. self.page.url = "https://auth.openai.com/sign-in-with-chatgpt/codex/consent"
  335. elif self.page.state == "consent":
  336. self.page.url = callback_url
  337. self.page.state = "callback"
  338. class _FakePage:
  339. def __init__(self):
  340. self.url = ""
  341. self.state = "new"
  342. self.filled = []
  343. self.clicked = []
  344. def goto(self, url, **_kwargs):
  345. self.url = url
  346. self.state = "email_prompt"
  347. def wait_for_timeout(self, *_args, **_kwargs):
  348. pass
  349. def locator(self, selector):
  350. return _FakeLocator(self, selector)
  351. page = _FakePage()
  352. captured = cpa_oauth.approve_openai_oauth_and_capture_callback(
  353. page,
  354. "https://auth.openai.com/oauth/authorize?client_id=codex",
  355. email_hint="n@example.com",
  356. log=lambda *_: None,
  357. timeout_sec=1,
  358. )
  359. self.assertEqual(captured, callback_url)
  360. self.assertEqual(page.filled, [("email_prompt", 'input[type="email"]', "n@example.com")])
  361. self.assertIn(("consent", 'button[type="submit"]'), page.clicked)
  362. def test_detects_callback_opened_in_another_context_page(self):
  363. import cpa_oauth
  364. callback_url = "http://localhost:1455/auth/callback?code=abc&state=xyz"
  365. class _FakeContext:
  366. def __init__(self):
  367. self.pages = []
  368. class _FakeCallbackPage:
  369. def __init__(self, url):
  370. self.url = url
  371. class _FakeLocator:
  372. def __init__(self, page, selector):
  373. self.page = page
  374. self.selector = selector
  375. def count(self):
  376. if 'button[type="submit"]' in self.selector:
  377. return 1
  378. return 0
  379. def nth(self, _idx):
  380. return self
  381. def is_visible(self):
  382. return True
  383. def is_enabled(self):
  384. return True
  385. def click(self, **_kwargs):
  386. self.page.clicked.append(self.selector)
  387. self.page.context.pages.append(_FakeCallbackPage(callback_url))
  388. class _FakePage:
  389. def __init__(self):
  390. self.url = ""
  391. self.clicked = []
  392. self.context = _FakeContext()
  393. self.context.pages.append(self)
  394. def goto(self, url, **_kwargs):
  395. self.url = "https://auth.openai.com/sign-in-with-chatgpt/codex/consent"
  396. def wait_for_timeout(self, *_args, **_kwargs):
  397. pass
  398. def locator(self, selector):
  399. return _FakeLocator(self, selector)
  400. page = _FakePage()
  401. captured = cpa_oauth.approve_openai_oauth_and_capture_callback(
  402. page,
  403. "https://auth.openai.com/oauth/authorize?client_id=codex",
  404. log=lambda *_: None,
  405. timeout_sec=1,
  406. )
  407. self.assertEqual(captured, callback_url)
  408. self.assertIn('button[type="submit"]', page.clicked)
  409. def test_capture_callback_even_when_localhost_navigation_raises(self):
  410. import cpa_oauth
  411. callback_url = "http://localhost:1455/auth/callback?code=abc&state=xyz"
  412. class _FakePage:
  413. def __init__(self):
  414. self.url = ""
  415. def goto(self, *_args, **_kwargs):
  416. self.url = callback_url
  417. raise RuntimeError("net::ERR_CONNECTION_REFUSED")
  418. def wait_for_timeout(self, *_args, **_kwargs):
  419. pass
  420. def locator(self, _selector):
  421. raise AssertionError("callback 已捕获时不应继续查找按钮")
  422. captured = cpa_oauth.approve_openai_oauth_and_capture_callback(
  423. _FakePage(),
  424. "https://auth.openai.com/oauth",
  425. log=lambda *_: None,
  426. )
  427. self.assertEqual(captured, callback_url)
  428. class AuthorizeFlowTests(unittest.TestCase):
  429. def test_authorize_keeps_cpa_panel_page_and_uses_separate_oauth_page(self):
  430. import cpa_oauth
  431. class _FakeContext:
  432. def __init__(self, oauth_page):
  433. self.oauth_page = oauth_page
  434. self.new_page_calls = 0
  435. def new_page(self):
  436. self.new_page_calls += 1
  437. return self.oauth_page
  438. class _FakePage:
  439. def __init__(self, name):
  440. self.name = name
  441. self.closed = False
  442. self.context = None
  443. def close(self):
  444. self.closed = True
  445. panel_page = _FakePage("panel")
  446. oauth_page = _FakePage("oauth")
  447. panel_page.context = _FakeContext(oauth_page)
  448. def fake_approve(page, oauth_url, **kwargs):
  449. self.assertIs(page, oauth_page)
  450. self.assertEqual(oauth_url, "https://auth.openai.com/oauth/authorize")
  451. self.assertEqual(kwargs["email_hint"], "n@example.com")
  452. return "http://localhost:1455/auth/callback?code=abc&state=xyz"
  453. def fake_submit(page, **kwargs):
  454. self.assertIs(page, panel_page)
  455. self.assertEqual(kwargs["callback_url"], "http://localhost:1455/auth/callback?code=abc&state=xyz")
  456. return "Authentication successful!"
  457. with patch.object(cpa_oauth, "fetch_cpa_oauth_url", return_value="https://auth.openai.com/oauth/authorize") as fetch:
  458. with patch.object(cpa_oauth, "approve_openai_oauth_and_capture_callback", side_effect=fake_approve) as approve:
  459. with patch.object(cpa_oauth, "submit_oauth_callback_to_cpa", side_effect=fake_submit) as submit:
  460. result = cpa_oauth.authorize_codex_oauth_to_cpa(
  461. panel_page,
  462. cpa_url="https://cpa.example/management.html#/oauth",
  463. management_key="secret",
  464. email_hint="n@example.com",
  465. log=lambda *_: None,
  466. )
  467. fetch.assert_called_once()
  468. approve.assert_called_once()
  469. submit.assert_called_once()
  470. self.assertEqual(panel_page.context.new_page_calls, 1)
  471. self.assertTrue(oauth_page.closed)
  472. self.assertFalse(panel_page.closed)
  473. self.assertEqual(result["callbackUrl"], "http://localhost:1455/auth/callback?code=abc&state=xyz")
  474. if __name__ == "__main__":
  475. unittest.main()