from __future__ import annotations import os import sys import unittest from contextlib import asynccontextmanager 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.playwright_worker import DouyinWorker class SecUserIdGuardTests(unittest.IsolatedAsyncioTestCase): async def test_im_direct_missing_sec_user_id_exits_before_online_state(self): worker = DouyinWorker(account_id=301, login_mode="im_direct") worker._require_sec_user_id = AsyncMock(return_value=False) 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.assertFalse(started) self.assertIn("sec_user_id", reason) worker._require_sec_user_id.assert_awaited_once() require_call = worker._require_sec_user_id.await_args self.assertTrue(require_call.kwargs["refresh_if_missing"]) self.assertTrue(require_call.kwargs["refresh_if_stale"]) worker._build_im_session_from_storage.assert_awaited_once() validate.assert_awaited_once_with(im_session) worker._persist_im_session.assert_not_awaited() worker._run_im_direct_service.assert_not_awaited() async def test_running_guard_precedes_disabled_follow_welcome_setting(self): worker = DouyinWorker(account_id=302, login_mode="im_direct") worker.is_running = True worker._im_service = SimpleNamespace(_running=True) worker._require_sec_user_id = AsyncMock(return_value=False) worker.get_db = AsyncMock() await worker.follow_welcome_tick() worker._require_sec_user_id.assert_awaited_once() require_call = worker._require_sec_user_id.await_args self.assertFalse(require_call.kwargs.get("refresh_if_missing", False)) # The identity guard must run before the account's follow-welcome flag # is queried. Otherwise accounts with that feature disabled could stay # hosted indefinitely without a sec_user_id. worker.get_db.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() 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_douyin_profile", 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_closes_browser_and_skips_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 False worker._load_user_agent = AsyncMock(return_value="test-agent") worker._probe_existing_login = AsyncMock(return_value=True) worker._finalize_login_session = AsyncMock() worker._require_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._require_sec_user_id.assert_awaited_once() require_call = worker._require_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() worker._persist_im_session.assert_not_awaited() worker._run_im_direct_service.assert_not_awaited() 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._require_sec_user_id = AsyncMock(return_value=True) 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._require_sec_user_id.assert_awaited_once() self.assertTrue(worker._require_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()