694 lines
26 KiB
Python
694 lines
26 KiB
Python
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._refresh_follow_welcome_config = AsyncMock()
|
|
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._refresh_follow_welcome_config.assert_not_awaited()
|
|
worker.get_db.assert_not_awaited()
|
|
|
|
async def test_cached_disabled_follow_setting_still_guards_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._require_sec_user_id = AsyncMock(return_value="")
|
|
|
|
await worker.follow_welcome_tick()
|
|
|
|
worker._refresh_follow_welcome_config.assert_awaited_once_with()
|
|
worker._require_sec_user_id.assert_awaited_once_with("托管运行中")
|
|
|
|
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()
|