from __future__ import annotations import json import os import sys import unittest from contextlib import asynccontextmanager from datetime import datetime, timedelta from pathlib import Path from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch BACKEND_DIR = Path(__file__).resolve().parents[1] os.environ["KEFU_DB_TYPE"] = "sqlite" os.environ["KEFU_DATABASE_URL"] = "" os.environ["KEFU_DB_PATH"] = str(BACKEND_DIR / "kefu.db") if str(BACKEND_DIR) not in sys.path: sys.path.insert(0, str(BACKEND_DIR)) from rpa_engine.douyin_im.session import DouyinImSession from rpa_engine import account_profile as account_profile_module from rpa_engine.credential import build_im_session_from_storage from rpa_engine.playwright_worker import DouyinWorker class SecUserIdGuardTests(unittest.IsolatedAsyncioTestCase): async def test_current_profile_uid_overrides_browser_runtime_uid(self): now = datetime.utcnow() row = SimpleNamespace( im_session_data=None, cookie_updated_at=now, uid="2609567359568155", profile_updated_at=now, ) result = MagicMock() result.one_or_none.return_value = row db = SimpleNamespace( execute=AsyncMock(return_value=result), close=AsyncMock(), ) worker = DouyinWorker(account_id=300, login_mode="im_direct") worker.get_db = AsyncMock(return_value=db) worker._load_raw_user_agent = AsyncMock(return_value="test-agent") storage = { "cookies": [{"name": "sessionid", "value": "test-session"}], "my_uid": 7678285795559818786, "origins": [ { "localStorage": [ { "name": "web_runtime_security_uid", "value": "7678285795559818786", } ] } ], } session = await worker._build_im_session_from_storage(storage) self.assertEqual(session.my_uid, 2609567359568155) # device_id 必须与权威 UID 同步（protobuf 发送时 device_id 优先取 # session.device_id，残留的浏览器采集值会导致 device_id != my_uid -> KICK） self.assertEqual(session.device_id, "2609567359568155") self.assertTrue(session.uid_verified) db.close.assert_awaited_once() async def test_stale_profile_uid_does_not_override_new_cookie(self): now = datetime.utcnow() row = SimpleNamespace( im_session_data=None, cookie_updated_at=now, uid="2609567359568155", profile_updated_at=now - timedelta(seconds=1), ) result = MagicMock() result.one_or_none.return_value = row db = SimpleNamespace( execute=AsyncMock(return_value=result), close=AsyncMock(), ) worker = DouyinWorker(account_id=300, login_mode="im_direct") worker.get_db = AsyncMock(return_value=db) worker._load_raw_user_agent = AsyncMock(return_value="test-agent") session = await worker._build_im_session_from_storage( { "cookies": [{"name": "sessionid", "value": "test-session"}], "my_uid": 7678285795559818786, } ) self.assertEqual(session.my_uid, 7678285795559818786) self.assertFalse(session.uid_verified) def test_persisted_verified_uid_wins_in_generic_session_builder(self): saved = DouyinImSession( cookies={"sessionid": "test-session"}, my_uid=2609567359568155, device_id="7678285795559818786", uid_verified=True, ) session = build_im_session_from_storage( { "cookies": [{"name": "sessionid", "value": "test-session"}], "my_uid": 7678285795559818786, "origins": [ { "localStorage": [ { "name": "web_runtime_security_uid", "value": "7678285795559818786", } ] } ], }, json.dumps(saved.to_dict()), ) self.assertEqual(session.my_uid, 2609567359568155) # device_id 同步为已核验 UID，避免凭证残留设备号导致 device_id != my_uid self.assertEqual(session.device_id, "2609567359568155") self.assertTrue(session.uid_verified) async def test_im_direct_missing_sec_user_id_continues_im_hosting(self): worker = DouyinWorker(account_id=301, login_mode="im_direct") worker._best_effort_sec_user_id = AsyncMock(return_value="") worker._load_user_agent = AsyncMock(return_value="test-agent") im_session = DouyinImSession( cookies={"sessionid": "test-session"}, my_uid=30101, ) worker._build_im_session_from_storage = AsyncMock(return_value=im_session) worker._persist_im_session = AsyncMock() worker._run_im_direct_service = AsyncMock() with patch( "rpa_engine.playwright_worker.validate_im_session", new=AsyncMock(return_value=(True, "ready")), ) as validate: started, reason = await worker._try_cookie_only_im_start( {"cookies": [{"name": "sessionid", "value": "test-session"}]} ) self.assertTrue(started) self.assertEqual(reason, "") worker._best_effort_sec_user_id.assert_awaited_once_with( refresh_if_missing=True, refresh_if_stale=True, ) worker._build_im_session_from_storage.assert_awaited_once() validate.assert_awaited_once_with(im_session) worker._persist_im_session.assert_awaited_once() worker._run_im_direct_service.assert_awaited_once_with(im_session) async def test_disabled_follow_welcome_does_not_require_identity(self): worker = DouyinWorker(account_id=302, login_mode="im_direct") worker.is_running = True worker._im_service = SimpleNamespace(_running=True) worker._best_effort_sec_user_id = AsyncMock() worker._refresh_follow_welcome_config = AsyncMock( return_value=(False, "", "") ) worker.get_db = AsyncMock() await worker.follow_welcome_tick() worker._refresh_follow_welcome_config.assert_awaited_once_with() worker._best_effort_sec_user_id.assert_not_awaited() worker.get_db.assert_not_awaited() async def test_cached_disabled_follow_setting_skips_missing_identity(self): worker = DouyinWorker(account_id=311, login_mode="im_direct") worker.is_running = True worker._im_service = SimpleNamespace(_running=True) worker._follow_config_loaded = True worker._refresh_follow_welcome_config = AsyncMock( return_value=(False, "", "") ) worker._best_effort_sec_user_id = AsyncMock(return_value="") await worker.follow_welcome_tick() worker._refresh_follow_welcome_config.assert_awaited_once_with() worker._best_effort_sec_user_id.assert_not_awaited() async def test_blank_sec_user_id_is_missing_and_stops_hosting(self): worker = DouyinWorker(account_id=303, login_mode="im_direct") worker._load_sec_user_id = AsyncMock(return_value=" ") worker._refresh_sec_user_id = AsyncMock(return_value="\t") worker._stop_for_missing_sec_user_id = AsyncMock() accepted = await worker._require_sec_user_id( "直连启动前", refresh_if_missing=True, ) self.assertFalse(accepted) worker._load_sec_user_id.assert_awaited() worker._refresh_sec_user_id.assert_awaited_once_with() worker._stop_for_missing_sec_user_id.assert_awaited_once_with("直连启动前") async def test_existing_sec_user_id_does_not_refresh_or_stop(self): worker = DouyinWorker(account_id=304, login_mode="im_direct") worker._load_sec_user_id = AsyncMock(return_value="MS4wLjABAAAA-valid") worker._refresh_sec_user_id = AsyncMock() worker._stop_for_missing_sec_user_id = AsyncMock() accepted = await worker._require_sec_user_id( "运行中", refresh_if_missing=True, ) self.assertTrue(accepted) worker._load_sec_user_id.assert_awaited_once_with() worker._refresh_sec_user_id.assert_not_awaited() worker._stop_for_missing_sec_user_id.assert_not_awaited() async def test_identity_read_error_does_not_claim_field_is_missing(self): worker = DouyinWorker(account_id=307, login_mode="im_direct") worker._load_sec_user_id = AsyncMock( side_effect=RuntimeError("database temporarily unavailable") ) worker._refresh_sec_user_id = AsyncMock() worker._stop_for_missing_sec_user_id = AsyncMock() with self.assertRaisesRegex(RuntimeError, "database temporarily unavailable"): await worker._require_sec_user_id( "托管运行中", refresh_if_missing=True, ) worker._refresh_sec_user_id.assert_not_awaited() worker._stop_for_missing_sec_user_id.assert_not_awaited() async def test_unknown_profile_refresh_error_does_not_stop_as_missing(self): worker = DouyinWorker(account_id=308, login_mode="im_direct") worker._load_sec_user_id = AsyncMock(return_value="") worker._refresh_sec_user_id = AsyncMock( side_effect=RuntimeError("暂时无法核验 sec_user_id，请稍后重试") ) worker._stop_for_missing_sec_user_id = AsyncMock() with self.assertRaisesRegex(RuntimeError, "暂时无法核验"): await worker._require_sec_user_id( "启动托管时", refresh_if_missing=True, ) worker._refresh_sec_user_id.assert_awaited_once_with() worker._stop_for_missing_sec_user_id.assert_not_awaited() async def test_stale_cookie_forces_refresh_instead_of_using_cached_identity(self): worker = DouyinWorker(account_id=309, login_mode="im_direct") worker._load_sec_user_id = AsyncMock(return_value="cached-sec-user-id") worker._sec_user_id_is_stale = AsyncMock(return_value=True) worker._refresh_sec_user_id = AsyncMock(return_value="fresh-sec-user-id") worker._stop_for_missing_sec_user_id = AsyncMock() resolved = await worker._require_sec_user_id( "启动托管时", refresh_if_missing=True, refresh_if_stale=True, ) self.assertEqual(resolved, "fresh-sec-user-id") worker._sec_user_id_is_stale.assert_awaited_once_with() worker._refresh_sec_user_id.assert_awaited_once_with() worker._stop_for_missing_sec_user_id.assert_not_awaited() async def test_profile_helper_uses_uid_fallback_for_sec_user_id(self): fetch_detail = AsyncMock( return_value={ "uid": "123456789", "nickname": "fallback-user", "sec_user_id": "", "fetched": True, "message": "", } ) run_in_thread = AsyncMock(return_value=("MS4wLjABAAAA-fallback", True)) with ( patch.object( account_profile_module, "fetch_douyin_profile_detail", fetch_detail, ), patch.object( account_profile_module.asyncio, "to_thread", run_in_thread, ), ): detail = await account_profile_module.fetch_douyin_profile_detail_with_sec_user_id( "cookie-json", "test-agent", ) self.assertEqual(detail["sec_user_id"], "MS4wLjABAAAA-fallback") self.assertEqual(detail["sec_user_id_status"], "found") fetch_detail.assert_awaited_once_with("cookie-json", "test-agent") run_in_thread.assert_awaited_once_with( account_profile_module._resolve_sec_user_id_by_uid_sync, "cookie-json", "test-agent", "123456789", ) async def test_profile_transport_failure_remains_unknown_not_missing(self): fetch_detail = AsyncMock( return_value={ "uid": "", "sec_user_id": "", "fetched": False, "message": "profile request timed out", } ) run_in_thread = AsyncMock() with ( patch.object( account_profile_module, "fetch_douyin_profile_detail", fetch_detail, ), patch.object( account_profile_module.asyncio, "to_thread", run_in_thread, ), ): detail = await account_profile_module.fetch_douyin_profile_detail_with_sec_user_id( "cookie-json", "test-agent", ) self.assertEqual(detail["sec_user_id_status"], "unknown") run_in_thread.assert_not_awaited() async def test_confirmed_empty_uid_fallback_is_missing(self): fetch_detail = AsyncMock( return_value={ "uid": "987654321", "sec_user_id": "", "fetched": True, "message": "", } ) run_in_thread = AsyncMock(return_value=("", True)) with ( patch.object( account_profile_module, "fetch_douyin_profile_detail", fetch_detail, ), patch.object( account_profile_module.asyncio, "to_thread", run_in_thread, ), ): detail = await account_profile_module.fetch_douyin_profile_detail_with_sec_user_id( "cookie-json", "test-agent", ) self.assertEqual(detail["sec_user_id"], "") self.assertEqual(detail["sec_user_id_status"], "missing") run_in_thread.assert_awaited_once() async def test_stop_for_missing_sec_user_id_is_offline_and_idempotent(self): worker = DouyinWorker(account_id=310, login_mode="im_direct") worker.is_running = True worker.stopping = False service = SimpleNamespace(_running=True) worker._im_service = service worker.update_account_status = AsyncMock() with patch( "rpa_engine.playwright_worker.system_logger.record" ) as record_system_log: await worker._stop_for_missing_sec_user_id("托管运行中") await worker._stop_for_missing_sec_user_id("重复检查") self.assertTrue(worker.stopping) self.assertFalse(worker.is_running) self.assertFalse(service._running) worker.update_account_status.assert_awaited_once() status_call = worker.update_account_status.await_args self.assertEqual(status_call.args[0], "offline") self.assertIn("sec_user_id", status_call.kwargs["error_msg"]) self.assertIn("托管已自动退出", status_call.kwargs["error_msg"]) record_system_log.assert_called_once() def test_profile_payload_uid_outranks_query_user_and_cookie_uid(self): """资料接口的 UID 必须压过 query/user 的 user_uid 和 cookie 兜底。 实测同一个 Cookie：query/user 返回 user_uid=938334054809296， 而账号资料 UID 是 2609567359568155。以前前者先占位，导致账号卡片 一直显示「用户938334054809296」，反查 sec_user_id 也永远失败。 """ auth = SimpleNamespace( cookie={}, msToken="test-ms-token", get_uid=lambda: "938334054809296", ) payloads = [ # 1) query/user：只有 user_uid，属于弱兜底 {"status_code": 0, "user_uid": "938334054809296"}, # 2) user/profile/self：权威账号资料 { "status_code": 0, "user": { "uid": "2609567359568155", "nickname": "凤的心", "sec_uid": "MS4wLjABAAAA-real-sec-user-id", "unique_id": "39688250979", "aweme_count": 50, "follower_count": 154, }, }, ] responses = [SimpleNamespace(json=lambda value=value: value) for value in payloads] with ( patch.object( account_profile_module, "_build_auth", return_value=(auth, "test-agent"), ), patch.object( account_profile_module.requests, "get", side_effect=responses, ), patch.object( account_profile_module, "generate_a_bogus", return_value="a-bogus", ), patch.object( account_profile_module, "generate_webid", return_value="web-id", ), patch.object( account_profile_module, "_requests_proxies", return_value=None, ), ): detail = account_profile_module.fetch_douyin_profile_detail_sync( "cookie-json", "test-agent", ) self.assertEqual(detail["uid"], "2609567359568155") self.assertEqual(detail["nickname"], "凤的心") self.assertTrue(detail["fetched"]) async def test_sync_writes_one_identity_to_both_account_and_profile(self): """卡片与详细资料必须来自同一次抓取，不能各抓一次抓出两个身份。""" detail = { "uid": "2609567359568155", "nickname": "凤的心", "avatar_url": "https://example.test/avatar.png", "unique_id": "39688250979", "signature": "", "sec_user_id": "MS4wLjABAAAA-real-sec-user-id", "sec_user_id_status": "found", "video_count": 50, "follower_count": 154, "following_count": 162, "total_favorited": 612, "favoriting_count": 0, "fetched": True, "message": "", } account = SimpleNamespace( id=1, username="用户938334054809296", douyin_uid="938334054809296", avatar_url=None, user_agent="test-agent", ) profile = SimpleNamespace( account_id=1, uid=None, nickname=None, avatar_url=None, unique_id=None, signature=None, sec_user_id=None, follower_count=None, following_count=None, total_favorited=None, favoriting_count=None, video_count=None, synced_at=None, sync_message=None, ) profile_result = MagicMock() profile_result.scalar_one_or_none.return_value = profile username_result = MagicMock() username_result.scalar_one_or_none.return_value = None db = SimpleNamespace( execute=AsyncMock(side_effect=lambda stmt: ( username_result if "accounts.username" in str(stmt).lower() or "username" in str(stmt).lower() else profile_result )), add=MagicMock(), commit=AsyncMock(), refresh=AsyncMock(), ) with ( patch.object( account_profile_module, "fetch_douyin_profile_detail_with_sec_user_id", new=AsyncMock(return_value=detail), ), patch.object( account_profile_module, "fetch_douyin_user_videos", new=AsyncMock(return_value={"videos": [], "message": "无作品"}), ), patch.object( account_profile_module, "fetch_douyin_profile", new=AsyncMock(side_effect=AssertionError("must not re-fetch")), ), patch.object( account_profile_module, "load_account_profile_from_db", new=AsyncMock(return_value={}), ), ): await account_profile_module.sync_account_profile_to_db( db, account, "cookie-json", ) # 同一个身份同时写进 accounts 行和 account_profile_details 行 self.assertEqual(account.douyin_uid, "2609567359568155") self.assertEqual(account.username, "凤的心") self.assertEqual(account.avatar_url, "https://example.test/avatar.png") self.assertEqual(profile.uid, "2609567359568155") self.assertEqual(profile.nickname, "凤的心") async def test_cookie_uid_without_valid_profile_payload_stays_unknown(self): auth = SimpleNamespace( cookie={}, msToken="test-ms-token", get_uid=lambda: "cookie-only-uid", ) payloads = [ {"status_code": 0, "data": {"status": "ok"}}, {"status_code": 0, "message": "success"}, {"status_code": 0, "data": {}}, ] responses = [SimpleNamespace(json=lambda value=value: value) for value in payloads] with ( patch.object( account_profile_module, "_build_auth", return_value=(auth, "test-agent"), ), patch.object( account_profile_module.requests, "get", side_effect=responses, ) as request_get, patch.object( account_profile_module, "generate_a_bogus", return_value="a-bogus", ), patch.object( account_profile_module, "generate_webid", return_value="web-id", ), patch.object( account_profile_module, "_requests_proxies", return_value=None, ), ): raw_detail = account_profile_module.fetch_douyin_profile_detail_sync( "cookie-json", "test-agent", ) self.assertEqual(request_get.call_count, 3) self.assertEqual(raw_detail["uid"], "cookie-only-uid") self.assertFalse(raw_detail["fetched"]) self.assertFalse(raw_detail["profile_response_valid"]) with patch.object( account_profile_module, "fetch_douyin_profile_detail", new=AsyncMock(return_value=raw_detail), ): checked = await account_profile_module.fetch_douyin_profile_detail_with_sec_user_id( "cookie-json", "test-agent", ) self.assertEqual(checked["sec_user_id_status"], "unknown") async def test_invalid_uid_fallback_payloads_stay_unknown(self): auth = SimpleNamespace(cookie={}, msToken="test-ms-token") cases = ( ( "nonzero-status", { "status_code": 10007, "data": { "user": { "uid": "fallback-uid", "sec_uid": "must-not-be-used", } }, }, ), ( "no-user-node", {"status_code": 0, "data": {"status": "ok"}}, ), ) for label, payload in cases: with self.subTest(label=label): with ( patch.object( account_profile_module, "_build_auth", return_value=(auth, "test-agent"), ), patch.object( account_profile_module, "_douyin_get_json", return_value=payload, ), patch.object( account_profile_module, "generate_webid", return_value="web-id", ), ): resolved, request_completed = ( account_profile_module._resolve_sec_user_id_by_uid_sync( "cookie-json", "test-agent", "fallback-uid", ) ) self.assertEqual(resolved, "") self.assertFalse(request_completed) fetch_detail = AsyncMock( return_value={ "uid": "fallback-uid", "sec_user_id": "", "fetched": True, "message": "", } ) run_in_thread = AsyncMock( return_value=(resolved, request_completed) ) with ( patch.object( account_profile_module, "fetch_douyin_profile_detail", fetch_detail, ), patch.object( account_profile_module.asyncio, "to_thread", run_in_thread, ), ): checked = await account_profile_module.fetch_douyin_profile_detail_with_sec_user_id( "cookie-json", "test-agent", ) self.assertEqual(checked["sec_user_id_status"], "unknown") async def test_unknown_profile_sync_preserves_cached_identity_and_videos(self): old_synced_at = object() profile = SimpleNamespace( sec_user_id="cached-sec-user-id", synced_at=old_synced_at, sync_message="previous message", ) account = SimpleNamespace( id=411, user_agent="test-agent", douyin_uid="cached-uid", username="cached-user", avatar_url="cached-avatar", ) scalar_result = SimpleNamespace(scalar_one_or_none=lambda: profile) db = SimpleNamespace( execute=AsyncMock(return_value=scalar_result), commit=AsyncMock(), add=MagicMock(), ) unknown_detail = { "uid": "cookie-uid", "sec_user_id": "", "sec_user_id_status": "unknown", "fetched": False, "message": "profile endpoint temporarily unavailable", } cached_response = { "account_id": 411, "sec_user_id": "cached-sec-user-id", "synced_at": "cached-time", "videos": [{"aweme_id": "cached-video"}], } fetch_videos = AsyncMock() apply_profile = AsyncMock() load_cached = AsyncMock(return_value=dict(cached_response)) with ( patch.object( account_profile_module, "fetch_douyin_profile_detail_with_sec_user_id", new=AsyncMock(return_value=unknown_detail), ), patch.object( account_profile_module, "fetch_douyin_user_videos", fetch_videos, ), patch.object( account_profile_module, "apply_profile_to_account", apply_profile, ), patch.object( account_profile_module, "load_account_profile_from_db", load_cached, ), ): result = await account_profile_module.sync_account_profile_to_db( db, account, "cookie-json", ) self.assertEqual(profile.sec_user_id, "cached-sec-user-id") self.assertIs(profile.synced_at, old_synced_at) fetch_videos.assert_not_awaited() apply_profile.assert_not_awaited() load_cached.assert_awaited_once_with(db, account) self.assertEqual(db.execute.await_count, 1) self.assertNotIn("delete", str(db.execute.await_args.args[0]).lower()) self.assertEqual(result["videos"], cached_response["videos"]) self.assertEqual( result["message"], "profile endpoint temporarily unavailable", ) async def test_browser_login_missing_sec_user_id_still_starts_im(self): worker = DouyinWorker(account_id=305, login_mode="browser") worker.is_running = True events: list[str] = [] async def reject_identity(*_args, **_kwargs): events.append("require-sec-user-id") return "" worker._load_user_agent = AsyncMock(return_value="test-agent") worker._probe_existing_login = AsyncMock(return_value=True) worker._finalize_login_session = AsyncMock() worker._best_effort_sec_user_id = AsyncMock(side_effect=reject_identity) worker._setup_im_network_listener = AsyncMock() worker._navigate_to_message_center = AsyncMock() worker._harvest_im_credentials = AsyncMock() worker._persist_cookies = AsyncMock() im_session = DouyinImSession( cookies={"sessionid": "test-session"}, my_uid=30501, ) worker._build_im_session = AsyncMock(return_value=im_session) worker._persist_im_session = AsyncMock() worker._run_im_direct_service = AsyncMock() worker._close_browser_only = AsyncMock() page = SimpleNamespace() context = SimpleNamespace( add_init_script=AsyncMock(), new_page=AsyncMock(return_value=page), ) browser = SimpleNamespace(new_context=AsyncMock(return_value=context)) playwright = SimpleNamespace() playwright_starter = SimpleNamespace( start=AsyncMock(return_value=playwright), ) @asynccontextmanager async def browser_slot(account_id: int, description: str): self.assertEqual(account_id, 305) self.assertEqual(description, "browser credential harvest") events.append("browser-acquired") try: yield finally: events.append("browser-released") controller = SimpleNamespace(browser_slot=browser_slot) with ( patch( "rpa_engine.playwright_worker.async_playwright", return_value=playwright_starter, ), patch( "rpa_engine.playwright_worker._launch_chromium", new=AsyncMock(return_value=browser), ), patch( "rpa_engine.playwright_worker.get_traffic_controller", return_value=controller, ), patch( "rpa_engine.playwright_worker.validate_im_session", new=AsyncMock(return_value=(True, "ready")), ), ): await worker._run_browser_im_flow( storage_state=None, cookie_info={ "cookie_valid": False, "has_sessionid": False, "reason": "no saved login", }, ) worker._finalize_login_session.assert_awaited_once_with() worker._best_effort_sec_user_id.assert_awaited_once() require_call = worker._best_effort_sec_user_id.await_args self.assertTrue(require_call.kwargs["force_refresh"]) self.assertEqual( events, ["browser-acquired", "browser-released", "require-sec-user-id"], ) worker._setup_im_network_listener.assert_awaited_once_with() worker._navigate_to_message_center.assert_awaited_once_with() worker._harvest_im_credentials.assert_awaited_once_with(timeout=25) worker._persist_cookies.assert_awaited_once_with() worker._build_im_session.assert_awaited_once_with() # sec_user_id 只服务关注欢迎语；缺它不能阻断私信托管， # 否则浏览器登录后账号立刻下线、永远不会自动回复。 worker._persist_im_session.assert_awaited_once() worker._run_im_direct_service.assert_awaited_once_with(im_session) worker._close_browser_only.assert_awaited_once_with() async def test_browser_login_with_sec_user_id_continues_to_im(self): worker = DouyinWorker(account_id=306, login_mode="browser") worker.is_running = True worker._load_user_agent = AsyncMock(return_value="test-agent") worker._probe_existing_login = AsyncMock(return_value=True) worker._finalize_login_session = AsyncMock() worker._best_effort_sec_user_id = AsyncMock(return_value="sec-uid-306") worker._setup_im_network_listener = AsyncMock() worker._navigate_to_message_center = AsyncMock() worker._harvest_im_credentials = AsyncMock() worker._persist_cookies = AsyncMock() im_session = DouyinImSession( cookies={"sessionid": "test-session"}, my_uid=30601, ) worker._build_im_session = AsyncMock(return_value=im_session) worker._persist_im_session = AsyncMock() worker._run_im_direct_service = AsyncMock() worker._close_browser_only = AsyncMock() page = SimpleNamespace() context = SimpleNamespace( add_init_script=AsyncMock(), new_page=AsyncMock(return_value=page), ) browser = SimpleNamespace(new_context=AsyncMock(return_value=context)) playwright_starter = SimpleNamespace( start=AsyncMock(return_value=SimpleNamespace()), ) @asynccontextmanager async def browser_slot(_account_id: int, _description: str): yield controller = SimpleNamespace(browser_slot=browser_slot) with ( patch( "rpa_engine.playwright_worker.async_playwright", return_value=playwright_starter, ), patch( "rpa_engine.playwright_worker._launch_chromium", new=AsyncMock(return_value=browser), ), patch( "rpa_engine.playwright_worker.get_traffic_controller", return_value=controller, ), patch( "rpa_engine.playwright_worker.validate_im_session", new=AsyncMock(return_value=(True, "ready")), ), ): await worker._run_browser_im_flow( storage_state=None, cookie_info={ "cookie_valid": False, "has_sessionid": False, "reason": "no saved login", }, ) worker._best_effort_sec_user_id.assert_awaited_once() self.assertTrue( worker._best_effort_sec_user_id.await_args.kwargs["force_refresh"] ) worker._setup_im_network_listener.assert_awaited_once_with() worker._navigate_to_message_center.assert_awaited_once_with() worker._harvest_im_credentials.assert_awaited_once_with(timeout=25) worker._build_im_session.assert_awaited_once_with() worker._persist_im_session.assert_awaited_once_with( im_session, status="online", clear_error=True, ) worker._run_im_direct_service.assert_awaited_once_with(im_session) worker._close_browser_only.assert_awaited() if __name__ == "__main__": unittest.main()