545 lines
20 KiB
Python
545 lines
20 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
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 import http_client as http_client_module
|
|
from rpa_engine.douyin_im.auth import DouyinAuth
|
|
from rpa_engine.douyin_im.frontier import ensure_frontier_ws
|
|
from rpa_engine.douyin_im.http_client import DouyinImHttpClient
|
|
from rpa_engine.douyin_im.session import DouyinImSession
|
|
from rpa_engine.egress_channels import EgressChannel
|
|
from rpa_engine import playwright_worker as playwright_worker_module
|
|
from rpa_engine.playwright_worker import DouyinWorker
|
|
|
|
|
|
class SendTextMessageEntryTests(unittest.IsolatedAsyncioTestCase):
|
|
def _make_client(self, account_id: int = 41) -> DouyinImHttpClient:
|
|
session = DouyinImSession(
|
|
cookies={"sessionid": "test-session"},
|
|
my_uid=10001,
|
|
conv_meta={"existing": {"ticket": "old-ticket"}},
|
|
)
|
|
return DouyinImHttpClient(session, account_id=account_id)
|
|
|
|
def test_query_user_does_not_replace_existing_im_uid(self):
|
|
client = self._make_client()
|
|
auth = SimpleNamespace(
|
|
source_ip="",
|
|
get_uid=MagicMock(return_value=938334054809296),
|
|
)
|
|
|
|
resolved = client._resolve_authoritative_uid(auth)
|
|
|
|
self.assertEqual(resolved, 10001)
|
|
self.assertEqual(client.session.my_uid, 10001)
|
|
auth.get_uid.assert_not_called()
|
|
|
|
def test_verified_uid_does_not_replace_the_runtime_device_id(self):
|
|
# device_id 是 query/user 返回的设备注册号，my_uid 是账号 UID。
|
|
# 拿 my_uid 顶替 device_id 会让 frontier 订阅到另一个地址：握手照样
|
|
# 成功，却永远收不到这个账号的私信。
|
|
session = DouyinImSession(
|
|
cookies={"sessionid": "test-session"},
|
|
my_uid=2609567359568155,
|
|
device_id="7678285795559818786",
|
|
web_id="7678286623234475535",
|
|
uid_verified=True,
|
|
)
|
|
|
|
auth = DouyinAuth.from_im_session(session)
|
|
|
|
self.assertEqual(auth.device_id, "7678285795559818786")
|
|
self.assertEqual(session.device_id, "7678285795559818786")
|
|
|
|
def test_captured_tokenless_browser_frontier_url_is_kept(self):
|
|
url = (
|
|
"wss://frontier100-normal.zijieapi.com/ws/v2?aid=6383&"
|
|
"device_platform=web&fpid=9&device_id=7678285795559818786&"
|
|
"access_key=0123456789abcdef0123456789abcdef&version_code=fws_1.0.0"
|
|
)
|
|
session = DouyinImSession(
|
|
cookies={"sessionid": "test-session"},
|
|
ws_urls=[url],
|
|
my_uid=2609567359568155,
|
|
device_id="7678285795559818786",
|
|
)
|
|
|
|
resolved = ensure_frontier_ws(session)
|
|
|
|
self.assertEqual(resolved, url)
|
|
self.assertEqual(session.ws_urls, [url])
|
|
|
|
def test_non_im_frontier_product_is_rejected(self):
|
|
url = (
|
|
"wss://frontier100-normal.zijieapi.com/ws/v2?aid=6383&"
|
|
"device_platform=web&fpid=971&device_id=7678285795559818786&"
|
|
"access_key=0123456789abcdef0123456789abcdef"
|
|
)
|
|
session = DouyinImSession(
|
|
cookies={"sessionid": "test-session"},
|
|
ws_urls=[url],
|
|
my_uid=2609567359568155,
|
|
device_id="7678285795559818786",
|
|
uid_verified=True,
|
|
)
|
|
|
|
resolved = ensure_frontier_ws(session)
|
|
|
|
self.assertIn("frontier-im.douyin.com", resolved)
|
|
self.assertIn("fpid=9", resolved)
|
|
self.assertIn("device_id=7678285795559818786", resolved)
|
|
self.assertNotIn("fpid=971", resolved)
|
|
|
|
async def test_public_entry_submits_the_whole_send_to_outbound_queue(self):
|
|
client = self._make_client(account_id=73)
|
|
submit = AsyncMock(return_value=True)
|
|
|
|
with patch(
|
|
"rpa_engine.douyin_im.traffic_control.submit_outbound",
|
|
submit,
|
|
):
|
|
sent = await client.send_text_message(
|
|
"0:1:10001:20002",
|
|
"hello",
|
|
conversation_short_id="short-1",
|
|
)
|
|
|
|
self.assertTrue(sent)
|
|
submit.assert_awaited_once()
|
|
account_id, operation = submit.await_args.args[:2]
|
|
self.assertEqual(account_id, 73)
|
|
self.assertTrue(callable(operation))
|
|
self.assertIn("20002", submit.await_args.kwargs["description"])
|
|
|
|
async def test_bypass_path_does_not_submit_again(self):
|
|
client = self._make_client()
|
|
client._resolve_authoritative_uid = MagicMock(return_value=0)
|
|
submit = AsyncMock(return_value=True)
|
|
|
|
with patch(
|
|
"rpa_engine.douyin_im.traffic_control.submit_outbound",
|
|
submit,
|
|
):
|
|
sent = await client.send_text_message(
|
|
"0:1:10001:20002",
|
|
"hello",
|
|
_bypass_global_queue=True,
|
|
)
|
|
|
|
self.assertFalse(sent)
|
|
submit.assert_not_awaited()
|
|
client._resolve_authoritative_uid.assert_called_once()
|
|
|
|
async def test_queued_client_result_state_is_copied_to_calling_client(self):
|
|
client = self._make_client(account_id=88)
|
|
queued_meta = {
|
|
"0:1:10001:20002": {
|
|
"ticket": "fresh-ticket",
|
|
"conversation_short_id": "fresh-short-id",
|
|
}
|
|
}
|
|
queued_client = SimpleNamespace(
|
|
send_text_message=AsyncMock(return_value=False),
|
|
last_send_meta=queued_meta,
|
|
last_error="credential expired",
|
|
last_send_needs_refresh=True,
|
|
last_send_channel_retryable=False,
|
|
last_request_debug="response status=401",
|
|
)
|
|
queued_context = MagicMock()
|
|
queued_context.__aenter__ = AsyncMock(return_value=queued_client)
|
|
queued_context.__aexit__ = AsyncMock(return_value=None)
|
|
queued_factory = MagicMock(return_value=queued_context)
|
|
|
|
async def execute_submission(account_id, operation, description=""):
|
|
self.assertEqual(account_id, 88)
|
|
self.assertIn("20002", description)
|
|
return await operation()
|
|
|
|
submit = AsyncMock(side_effect=execute_submission)
|
|
|
|
with (
|
|
patch(
|
|
"rpa_engine.douyin_im.traffic_control.submit_outbound",
|
|
submit,
|
|
),
|
|
patch.object(
|
|
http_client_module,
|
|
"DouyinImHttpClient",
|
|
queued_factory,
|
|
),
|
|
):
|
|
sent = await client.send_text_message(
|
|
"0:1:10001:20002",
|
|
"queued hello",
|
|
conversation_short_id="short-before-send",
|
|
)
|
|
|
|
self.assertFalse(sent)
|
|
submit.assert_awaited_once()
|
|
queued_factory.assert_called_once_with(client.session, account_id=88)
|
|
queued_client.send_text_message.assert_awaited_once_with(
|
|
"0:1:10001:20002",
|
|
"queued hello",
|
|
conversation_short_id="short-before-send",
|
|
_bypass_global_queue=True,
|
|
)
|
|
self.assertEqual(client.last_send_meta, queued_meta)
|
|
self.assertIsNot(client.last_send_meta, queued_meta)
|
|
self.assertEqual(client.last_error, "credential expired")
|
|
self.assertTrue(client.last_send_needs_refresh)
|
|
self.assertEqual(client.last_request_debug, "response status=401")
|
|
|
|
async def test_retryable_network_failure_switches_channels_serially(self):
|
|
client = self._make_client(account_id=89)
|
|
client.session.egress_auto_attempts = 2
|
|
routes = [
|
|
EgressChannel("198.51.100.10", "10.0.0.10", "eth0", True),
|
|
EgressChannel("198.51.100.11", "10.0.0.11", "eth0:1", False),
|
|
]
|
|
|
|
first = SimpleNamespace(
|
|
send_text_message=AsyncMock(return_value=False),
|
|
last_send_meta={},
|
|
last_error="connect timeout",
|
|
last_send_needs_refresh=False,
|
|
last_send_channel_retryable=True,
|
|
last_request_debug="first route",
|
|
)
|
|
second = SimpleNamespace(
|
|
send_text_message=AsyncMock(return_value=True),
|
|
last_send_meta={"conv": {"ticket": "ok"}},
|
|
last_error="",
|
|
last_send_needs_refresh=False,
|
|
last_send_channel_retryable=False,
|
|
last_request_debug="second route",
|
|
)
|
|
|
|
def context_for(value):
|
|
context = MagicMock()
|
|
context.__aenter__ = AsyncMock(return_value=value)
|
|
context.__aexit__ = AsyncMock(return_value=None)
|
|
return context
|
|
|
|
queued_factory = MagicMock(side_effect=[context_for(first), context_for(second)])
|
|
|
|
async def execute_submission(account_id, operation, description=""):
|
|
self.assertEqual(account_id, 89)
|
|
return await operation()
|
|
|
|
with (
|
|
patch(
|
|
"rpa_engine.douyin_im.traffic_control.submit_outbound",
|
|
AsyncMock(side_effect=execute_submission),
|
|
),
|
|
patch(
|
|
"rpa_engine.douyin_im.http_client.resolve_send_channels",
|
|
AsyncMock(return_value=routes),
|
|
),
|
|
patch.object(http_client_module, "DouyinImHttpClient", queued_factory),
|
|
patch.object(http_client_module.system_logger, "record"),
|
|
):
|
|
sent = await client.send_text_message("0:1:10001:20002", "hello")
|
|
|
|
self.assertTrue(sent)
|
|
self.assertEqual(queued_factory.call_count, 2)
|
|
self.assertEqual(queued_factory.call_args_list[0].kwargs["source_ip"], "10.0.0.10")
|
|
self.assertEqual(queued_factory.call_args_list[1].kwargs["source_ip"], "10.0.0.11")
|
|
first.send_text_message.assert_awaited_once()
|
|
second.send_text_message.assert_awaited_once()
|
|
self.assertEqual(client.last_request_debug, "second route")
|
|
|
|
async def test_kick_never_switches_public_channels(self):
|
|
client = self._make_client(account_id=90)
|
|
client.session.egress_auto_attempts = 2
|
|
routes = [
|
|
EgressChannel("198.51.100.10", "10.0.0.10", "eth0", True),
|
|
EgressChannel("198.51.100.11", "10.0.0.11", "eth1", False),
|
|
]
|
|
kicked = SimpleNamespace(
|
|
send_text_message=AsyncMock(return_value=False),
|
|
last_send_meta={},
|
|
last_error="decision=KICK",
|
|
last_send_needs_refresh=False,
|
|
last_send_channel_retryable=False,
|
|
last_request_debug="terminal kick",
|
|
)
|
|
context = MagicMock()
|
|
context.__aenter__ = AsyncMock(return_value=kicked)
|
|
context.__aexit__ = AsyncMock(return_value=None)
|
|
queued_factory = MagicMock(return_value=context)
|
|
|
|
async def execute_submission(account_id, operation, description=""):
|
|
return await operation()
|
|
|
|
with (
|
|
patch(
|
|
"rpa_engine.douyin_im.traffic_control.submit_outbound",
|
|
AsyncMock(side_effect=execute_submission),
|
|
),
|
|
patch(
|
|
"rpa_engine.douyin_im.http_client.resolve_send_channels",
|
|
AsyncMock(return_value=routes),
|
|
),
|
|
patch.object(http_client_module, "DouyinImHttpClient", queued_factory),
|
|
):
|
|
sent = await client.send_text_message("0:1:10001:20002", "hello")
|
|
|
|
self.assertFalse(sent)
|
|
queued_factory.assert_called_once()
|
|
kicked.send_text_message.assert_awaited_once()
|
|
|
|
|
|
class WorkerLifecycleTests(unittest.IsolatedAsyncioTestCase):
|
|
async def test_browser_launch_uses_selected_source_proxy(self):
|
|
launch = AsyncMock(return_value="browser")
|
|
pw = SimpleNamespace(chromium=SimpleNamespace(launch=launch))
|
|
source_proxy = AsyncMock(return_value={"server": "http://127.0.0.1:43210"})
|
|
|
|
with (
|
|
patch.object(
|
|
playwright_worker_module,
|
|
"ensure_browser_display",
|
|
AsyncMock(),
|
|
),
|
|
patch.object(
|
|
playwright_worker_module,
|
|
"playwright_proxy_for_source",
|
|
source_proxy,
|
|
),
|
|
patch.object(playwright_worker_module, "playwright_proxy") as global_proxy,
|
|
):
|
|
browser = await playwright_worker_module._launch_chromium(
|
|
pw,
|
|
["--no-sandbox"],
|
|
headless=True,
|
|
source_ip="10.0.0.6",
|
|
)
|
|
|
|
self.assertEqual(browser, "browser")
|
|
source_proxy.assert_awaited_once_with("10.0.0.6")
|
|
global_proxy.assert_not_called()
|
|
launch.assert_awaited_once_with(
|
|
headless=True,
|
|
args=["--no-sandbox"],
|
|
proxy={"server": "http://127.0.0.1:43210"},
|
|
)
|
|
|
|
async def test_virtual_display_starts_before_playwright_driver(self):
|
|
ensure_display = AsyncMock()
|
|
|
|
async def start_driver():
|
|
ensure_display.assert_awaited_once_with(False)
|
|
return "playwright-driver"
|
|
|
|
manager = SimpleNamespace(start=AsyncMock(side_effect=start_driver))
|
|
with (
|
|
patch.object(
|
|
playwright_worker_module,
|
|
"resolve_headless",
|
|
return_value=False,
|
|
),
|
|
patch.object(
|
|
playwright_worker_module,
|
|
"ensure_browser_display",
|
|
ensure_display,
|
|
),
|
|
patch.object(
|
|
playwright_worker_module,
|
|
"async_playwright",
|
|
return_value=manager,
|
|
),
|
|
):
|
|
driver, headless = await playwright_worker_module._start_playwright_for_browser()
|
|
|
|
self.assertEqual(driver, "playwright-driver")
|
|
self.assertFalse(headless)
|
|
manager.start.assert_awaited_once_with()
|
|
|
|
async def test_visible_login_prompt_overrides_stale_sessionid(self):
|
|
worker = DouyinWorker(account_id=917, login_mode="browser")
|
|
worker._has_visible_login_prompt = AsyncMock(return_value=True)
|
|
worker.check_homepage_login_status = AsyncMock(return_value=True)
|
|
worker.check_logged_in_by_cookie = AsyncMock(return_value=True)
|
|
|
|
logged_in = await worker._verify_login_state()
|
|
|
|
self.assertFalse(logged_in)
|
|
worker.check_homepage_login_status.assert_not_awaited()
|
|
worker.check_logged_in_by_cookie.assert_not_awaited()
|
|
|
|
async def test_homepage_message_entry_is_not_login_evidence(self):
|
|
worker = DouyinWorker(account_id=918, login_mode="browser")
|
|
worker._has_visible_login_prompt = AsyncMock(return_value=False)
|
|
worker.page = SimpleNamespace(query_selector=AsyncMock(return_value=None))
|
|
|
|
logged_in = await worker.check_homepage_login_status()
|
|
|
|
self.assertFalse(logged_in)
|
|
worker.page.query_selector.assert_awaited_once()
|
|
|
|
async def test_start_saves_task_and_stop_waits_until_it_is_done(self):
|
|
worker = DouyinWorker(account_id=919, login_mode="im_direct")
|
|
loop_started = asyncio.Event()
|
|
loop_finished = asyncio.Event()
|
|
never_finishes_without_cancellation = asyncio.Event()
|
|
|
|
async def controlled_run_loop() -> None:
|
|
loop_started.set()
|
|
try:
|
|
await never_finishes_without_cancellation.wait()
|
|
finally:
|
|
loop_finished.set()
|
|
|
|
worker._run_loop = controlled_run_loop
|
|
worker.update_account_status = AsyncMock()
|
|
worker.cleanup = AsyncMock()
|
|
|
|
await worker.start()
|
|
task = worker._task
|
|
|
|
self.assertIsInstance(task, asyncio.Task)
|
|
self.assertEqual(task.get_name(), "douyin-worker-919")
|
|
self.assertFalse(task.done())
|
|
await asyncio.wait_for(loop_started.wait(), timeout=0.2)
|
|
|
|
# Starting an already-running worker must retain the same owned task.
|
|
await worker.start()
|
|
self.assertIs(worker._task, task)
|
|
|
|
await asyncio.wait_for(worker.stop(), timeout=0.5)
|
|
|
|
self.assertTrue(task.done())
|
|
self.assertTrue(task.cancelled())
|
|
self.assertTrue(loop_finished.is_set())
|
|
self.assertIsNone(worker._task)
|
|
self.assertTrue(worker.stopping)
|
|
self.assertFalse(worker.is_running)
|
|
worker.update_account_status.assert_awaited_once_with("offline")
|
|
worker.cleanup.assert_awaited_once_with()
|
|
leaked = [
|
|
running
|
|
for running in asyncio.all_tasks()
|
|
if running is not asyncio.current_task()
|
|
and running.get_name() == "douyin-worker-919"
|
|
and not running.done()
|
|
]
|
|
self.assertEqual(leaked, [])
|
|
|
|
async def test_browser_flow_closes_resources_before_releasing_slot_on_error(self):
|
|
worker = DouyinWorker(account_id=920, login_mode="browser")
|
|
events: list[str] = []
|
|
|
|
# A truthy browser resource makes the flow's finally block responsible
|
|
# for cleanup even though credential preparation fails midway.
|
|
worker.page = object()
|
|
|
|
async def failing_prepare(storage_state, cookie_info):
|
|
events.append("prepare")
|
|
raise RuntimeError("credential harvest failed")
|
|
|
|
async def close_browser_only():
|
|
events.append("close")
|
|
worker.page = None
|
|
|
|
@asynccontextmanager
|
|
async def browser_slot(account_id, description):
|
|
self.assertEqual(account_id, 920)
|
|
self.assertEqual(description, "browser credential harvest")
|
|
events.append("acquire")
|
|
try:
|
|
yield
|
|
finally:
|
|
events.append("release")
|
|
|
|
controller = SimpleNamespace(browser_slot=browser_slot)
|
|
worker._prepare_browser_im_flow = AsyncMock(side_effect=failing_prepare)
|
|
worker._close_browser_only = AsyncMock(side_effect=close_browser_only)
|
|
|
|
with patch(
|
|
"rpa_engine.playwright_worker.get_traffic_controller",
|
|
return_value=controller,
|
|
):
|
|
with self.assertRaisesRegex(RuntimeError, "credential harvest failed"):
|
|
await worker._run_browser_im_flow(
|
|
storage_state={"cookies": []},
|
|
cookie_info={"has_sessionid": True},
|
|
)
|
|
|
|
self.assertEqual(events, ["acquire", "prepare", "close", "release"])
|
|
worker._prepare_browser_im_flow.assert_awaited_once_with(
|
|
{"cookies": []},
|
|
{"has_sessionid": True},
|
|
)
|
|
worker._close_browser_only.assert_awaited_once_with()
|
|
|
|
async def test_close_browser_only_closes_every_resource_and_clears_references(self):
|
|
worker = DouyinWorker(account_id=921, login_mode="browser")
|
|
events: list[str] = []
|
|
|
|
async def record(name: str, *, fail: bool = False) -> None:
|
|
events.append(name)
|
|
if fail:
|
|
raise RuntimeError(f"{name} close failed")
|
|
|
|
async def close_page() -> None:
|
|
await record("page")
|
|
|
|
async def close_context() -> None:
|
|
await record("context", fail=True)
|
|
|
|
async def close_browser() -> None:
|
|
await record("browser")
|
|
|
|
async def stop_playwright() -> None:
|
|
await record("playwright")
|
|
|
|
page = SimpleNamespace(
|
|
close=AsyncMock(side_effect=close_page),
|
|
)
|
|
context = SimpleNamespace(
|
|
# One failed close must not keep the remaining processes alive.
|
|
close=AsyncMock(side_effect=close_context),
|
|
)
|
|
browser = SimpleNamespace(
|
|
close=AsyncMock(side_effect=close_browser),
|
|
)
|
|
playwright = SimpleNamespace(
|
|
stop=AsyncMock(side_effect=stop_playwright),
|
|
)
|
|
worker.page = page
|
|
worker.context = context
|
|
worker.browser = browser
|
|
worker.playwright = playwright
|
|
|
|
await worker._close_browser_only()
|
|
|
|
self.assertEqual(events, ["page", "context", "browser", "playwright"])
|
|
page.close.assert_awaited_once_with()
|
|
context.close.assert_awaited_once_with()
|
|
browser.close.assert_awaited_once_with()
|
|
playwright.stop.assert_awaited_once_with()
|
|
self.assertIsNone(worker.page)
|
|
self.assertIsNone(worker.context)
|
|
self.assertIsNone(worker.browser)
|
|
self.assertIsNone(worker.playwright)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|