| 1 | import asyncio |
| 2 | import time |
| 3 | |
| 4 | import pytest |
| 5 | |
| 6 | from tools.cookie_fetcher import ( |
| 7 | extract_ms_token_from_text, |
| 8 | filter_cookies, |
| 9 | goto_with_fallback, |
| 10 | try_extract_ms_token, |
| 11 | wait_for_login_confirmation, |
| 12 | ) |
| 13 | |
| 14 | |
| 15 | class FakePage: |
| 16 | def __init__(self, outcomes): |
| 17 | self._outcomes = list(outcomes) |
| 18 | self.calls = [] |
| 19 | |
| 20 | async def goto(self, url, wait_until=None, timeout=None): |
| 21 | self.calls.append( |
| 22 | { |
| 23 | "url": url, |
| 24 | "wait_until": wait_until, |
| 25 | "timeout": timeout, |
| 26 | } |
| 27 | ) |
| 28 | outcome = self._outcomes.pop(0) |
| 29 | if isinstance(outcome, Exception): |
| 30 | raise outcome |
| 31 | return outcome |
| 32 | |
| 33 | |
| 34 | class SlowPage: |
| 35 | def __init__(self): |
| 36 | self.calls = [] |
| 37 | self.cancelled = False |
| 38 | |
| 39 | async def goto(self, url, wait_until=None, timeout=None): |
| 40 | self.calls.append( |
| 41 | { |
| 42 | "url": url, |
| 43 | "wait_until": wait_until, |
| 44 | "timeout": timeout, |
| 45 | } |
| 46 | ) |
| 47 | try: |
| 48 | await asyncio.sleep(60) |
| 49 | except asyncio.CancelledError: |
| 50 | self.cancelled = True |
| 51 | raise |
| 52 | |
| 53 | async def evaluate(self, _): |
| 54 | return "" |
| 55 | |
| 56 | |
| 57 | def test_goto_with_fallback_when_networkidle_timeout(): |
| 58 | page = FakePage([TimeoutError("network idle timeout"), object()]) |
| 59 | wait_until = asyncio.run(goto_with_fallback(page, "https://www.douyin.com/")) |
| 60 | |
| 61 | assert wait_until == "domcontentloaded" |
| 62 | assert page.calls[0]["wait_until"] == "networkidle" |
| 63 | assert page.calls[1]["wait_until"] == "domcontentloaded" |
| 64 | |
| 65 | |
| 66 | def test_goto_with_fallback_raises_non_timeout_errors(): |
| 67 | page = FakePage([RuntimeError("unexpected error")]) |
| 68 | |
| 69 | with pytest.raises(RuntimeError, match="unexpected error"): |
| 70 | asyncio.run(goto_with_fallback(page, "https://www.douyin.com/")) |
| 71 | |
| 72 | assert len(page.calls) == 1 |
| 73 | |
| 74 | |
| 75 | def test_goto_with_fallback_handles_target_closed(): |
| 76 | class TargetClosedError(Exception): |
| 77 | pass |
| 78 | |
| 79 | page = FakePage([TargetClosedError("Target page, context or browser has been closed")]) |
| 80 | wait_until = asyncio.run(goto_with_fallback(page, "https://www.douyin.com/")) |
| 81 | |
| 82 | assert wait_until == "target_closed" |
| 83 | assert len(page.calls) == 1 |
| 84 | |
| 85 | |
| 86 | def test_goto_with_fallback_returns_timeout_when_fallback_also_times_out(): |
| 87 | page = FakePage([TimeoutError("primary timeout"), TimeoutError("fallback timeout")]) |
| 88 | wait_until = asyncio.run(goto_with_fallback(page, "https://www.douyin.com/")) |
| 89 | |
| 90 | assert wait_until == "timeout" |
| 91 | assert len(page.calls) == 2 |
| 92 | |
| 93 | |
| 94 | def test_wait_for_login_confirmation_returns_without_waiting_navigation(): |
| 95 | page = SlowPage() |
| 96 | started = time.time() |
| 97 | |
| 98 | asyncio.run( |
| 99 | wait_for_login_confirmation( |
| 100 | page, |
| 101 | "https://www.douyin.com/", |
| 102 | input_func=lambda: "", |
| 103 | ) |
| 104 | ) |
| 105 | elapsed = time.time() - started |
| 106 | |
| 107 | assert elapsed < 1 |
| 108 | assert len(page.calls) == 1 |
| 109 | assert page.cancelled is True |
| 110 | |
| 111 | |
| 112 | def test_wait_for_login_confirmation_handles_completed_navigation(): |
| 113 | page = FakePage([object()]) |
| 114 | |
| 115 | asyncio.run( |
| 116 | wait_for_login_confirmation( |
| 117 | page, |
| 118 | "https://www.douyin.com/", |
| 119 | input_func=lambda: "", |
| 120 | ) |
| 121 | ) |
| 122 | |
| 123 | assert len(page.calls) == 1 |
| 124 | |
| 125 | |
| 126 | def test_try_extract_ms_token_from_observed_headers(): |
| 127 | page = SlowPage() |
| 128 | |
| 129 | token = asyncio.run( |
| 130 | try_extract_ms_token( |
| 131 | page, |
| 132 | {"ttwid": "x"}, |
| 133 | ["ttwid=abc; msToken=token-from-header"], |
| 134 | [], |
| 135 | ) |
| 136 | ) |
| 137 | |
| 138 | assert token == "token-from-header" |
| 139 | |
| 140 | |
| 141 | def test_extract_ms_token_from_text_supports_json_and_query_formats(): |
| 142 | assert ( |
| 143 | extract_ms_token_from_text("https://www.douyin.com/?foo=1&msToken=query-token&bar=2") |
| 144 | == "query-token" |
| 145 | ) |
| 146 | assert extract_ms_token_from_text('{"msToken":"json-token","x":1}') == "json-token" |
| 147 | |
| 148 | |
| 149 | def test_filter_cookies_keeps_waf_and_fingerprint_keys_but_drops_unrelated_keys(): |
| 150 | cookies = filter_cookies( |
| 151 | { |
| 152 | "ttwid": "ttwid-token", |
| 153 | "msToken": "ms-token", |
| 154 | "_waftokenid": "waf-token", |
| 155 | "s_v_web_id": "verify-id", |
| 156 | "__ac_signature": "ac-signature", |
| 157 | "random_cookie": "should-be-filtered", |
| 158 | } |
| 159 | ) |
| 160 | |
| 161 | assert cookies["ttwid"] == "ttwid-token" |
| 162 | assert cookies["msToken"] == "ms-token" |
| 163 | assert cookies["_waftokenid"] == "waf-token" |
| 164 | assert cookies["s_v_web_id"] == "verify-id" |
| 165 | assert cookies["__ac_signature"] == "ac-signature" |
| 166 | assert "random_cookie" not in cookies |
| 167 |