444 lines
16 KiB
Python
444 lines
16 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, 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.http_client import DouyinImHttpClient
|
|
from rpa_engine.douyin_im.session import DouyinImSession
|
|
from rpa_engine.douyin_im.service import DouyinImService, _conversation_poll_timing
|
|
from rpa_engine.douyin_im import service as service_module
|
|
|
|
|
|
class ConversationPollBandwidthTests(unittest.IsolatedAsyncioTestCase):
|
|
def _make_client(self) -> DouyinImHttpClient:
|
|
return DouyinImHttpClient(
|
|
DouyinImSession(cookies={"sessionid": "test"}, my_uid=10001),
|
|
account_id=9,
|
|
)
|
|
|
|
async def test_inbox_is_fetched_with_exactly_one_protobuf_request(self):
|
|
"""imapi 只认 protobuf；轮询一次就只该发一个请求。"""
|
|
client = self._make_client()
|
|
client.fetch_inbox_messages = AsyncMock(return_value=[])
|
|
client._request = AsyncMock()
|
|
|
|
self.assertEqual(await client.get_conversations(), [])
|
|
client.fetch_inbox_messages.assert_awaited_once()
|
|
# 不能再退回 JSON 的 /v1/conversation/list：那个请求恒被抖音拒绝。
|
|
client._request.assert_not_awaited()
|
|
|
|
async def test_transport_failure_is_recorded_and_returns_empty(self):
|
|
client = self._make_client()
|
|
client.fetch_inbox_messages = AsyncMock(side_effect=RuntimeError("boom"))
|
|
|
|
with self.assertLogs("douyin_im.http", level="WARNING"):
|
|
self.assertEqual(await client.get_conversations(), [])
|
|
|
|
self.assertIn("boom", client.last_error)
|
|
|
|
async def test_inbox_messages_group_into_one_row_per_conversation(self):
|
|
client = self._make_client()
|
|
client.fetch_inbox_messages = AsyncMock(
|
|
return_value=[
|
|
{
|
|
"conversation_id": "0:1:10001:20001",
|
|
"server_message_id": "700",
|
|
"conversation_short_id": "555",
|
|
"message_type": 7,
|
|
"sender": "20001",
|
|
"content": '{"text":"旧"}',
|
|
},
|
|
{
|
|
"conversation_id": "0:1:10001:20001",
|
|
"server_message_id": "900",
|
|
"conversation_short_id": "555",
|
|
"message_type": 7,
|
|
"sender": "20001",
|
|
"content": '{"text":"新"}',
|
|
},
|
|
{
|
|
"conversation_id": "0:1:10001:20002",
|
|
"server_message_id": "800",
|
|
"message_type": 7,
|
|
"sender": "20002",
|
|
"content": '{"text":"另一个"}',
|
|
},
|
|
]
|
|
)
|
|
|
|
rows = await client.get_conversations(enrich_profiles=False)
|
|
|
|
by_id = {r["conversation_id"]: r for r in rows}
|
|
self.assertEqual(len(rows), 2)
|
|
# 同一会话只保留 server_message_id 最大的那条
|
|
self.assertEqual(by_id["0:1:10001:20001"]["server_message_id"], "900")
|
|
self.assertIn("新", by_id["0:1:10001:20001"]["content"])
|
|
# peer_uid 由 conversation_id 推导，不能直接取 sender（可能是自己）
|
|
self.assertEqual(by_id["0:1:10001:20002"]["peer_uid"], "20002")
|
|
# 顺手缓存 short_id，发送时就不必再 create 一次会话
|
|
self.assertEqual(
|
|
client.session.conv_meta["0:1:10001:20001"]["conversation_short_id"],
|
|
"555",
|
|
)
|
|
|
|
def test_websocket_reconciliation_is_slow_and_http_fallback_stays_fast(self):
|
|
with patch.dict(
|
|
os.environ,
|
|
{
|
|
"KEFU_WS_RECONCILE_INTERVAL_SECONDS": "120",
|
|
"KEFU_HTTP_POLL_INTERVAL_SECONDS": "15",
|
|
},
|
|
):
|
|
ws_interval, ws_stagger = _conversation_poll_timing(123, True)
|
|
http_interval, http_stagger = _conversation_poll_timing(123, False)
|
|
|
|
self.assertEqual(ws_interval, 120)
|
|
self.assertEqual(http_interval, 15)
|
|
self.assertGreaterEqual(ws_stagger, 0)
|
|
self.assertLess(ws_stagger, ws_interval)
|
|
self.assertGreaterEqual(http_stagger, 0)
|
|
self.assertLess(http_stagger, http_interval)
|
|
|
|
def test_poll_interval_expands_to_the_configured_population_budget(self):
|
|
with patch.dict(
|
|
os.environ,
|
|
{
|
|
"KEFU_WS_RECONCILE_INTERVAL_SECONDS": "120",
|
|
"KEFU_HTTP_POLL_INTERVAL_SECONDS": "15",
|
|
"KEFU_WS_POLL_BUDGET_RPS": "1",
|
|
"KEFU_HTTP_POLL_BUDGET_RPS": "1",
|
|
},
|
|
):
|
|
ws_interval, _ = _conversation_poll_timing(
|
|
123,
|
|
True,
|
|
population=500,
|
|
)
|
|
http_interval, _ = _conversation_poll_timing(
|
|
123,
|
|
False,
|
|
population=500,
|
|
)
|
|
|
|
self.assertEqual(ws_interval, 500)
|
|
self.assertEqual(http_interval, 500)
|
|
|
|
def test_initial_unread_concurrency_is_configurable_and_bounded(self):
|
|
with patch.dict(
|
|
os.environ,
|
|
{"KEFU_INITIAL_UNREAD_CONCURRENCY": "4"},
|
|
):
|
|
self.assertEqual(service_module._initial_unread_concurrency(), 4)
|
|
with patch.dict(
|
|
os.environ,
|
|
{"KEFU_INITIAL_UNREAD_CONCURRENCY": "999"},
|
|
):
|
|
self.assertEqual(service_module._initial_unread_concurrency(), 8)
|
|
|
|
async def test_service_poll_uses_one_conversation_request_without_unread_probe(self):
|
|
class _Controller:
|
|
@asynccontextmanager
|
|
async def background_slot(self, *_args, **_kwargs):
|
|
yield
|
|
|
|
class _HttpClient:
|
|
def __init__(self):
|
|
self.get_conversations = AsyncMock(return_value=[])
|
|
self.conversation_list_unsupported = False
|
|
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, *_args):
|
|
return False
|
|
|
|
http = _HttpClient()
|
|
service = DouyinImService(
|
|
session=DouyinImSession(cookies={"sessionid": "test"}, my_uid=10001),
|
|
match_reply=AsyncMock(return_value=None),
|
|
log_fn=AsyncMock(),
|
|
account_id=9,
|
|
)
|
|
service._index_conversations = AsyncMock()
|
|
|
|
with (
|
|
patch.object(service_module, "get_traffic_controller", return_value=_Controller()),
|
|
patch.object(service_module, "DouyinImHttpClient", return_value=http),
|
|
):
|
|
await service._poll_conversations()
|
|
|
|
http.get_conversations.assert_awaited_once_with(enrich_profiles=False)
|
|
service._index_conversations.assert_awaited_once_with(
|
|
[],
|
|
enrich_profiles=False,
|
|
)
|
|
|
|
async def test_poll_only_handles_unread_or_a_genuinely_changed_preview(self):
|
|
class _Controller:
|
|
def __init__(self):
|
|
self.startup_flags = []
|
|
|
|
@asynccontextmanager
|
|
async def background_slot(self, *_args, **kwargs):
|
|
self.startup_flags.append(bool(kwargs.get("startup")))
|
|
yield
|
|
|
|
snapshots = [
|
|
[
|
|
{
|
|
"conversation_id": "0:1:10001:20001",
|
|
"peer_uid": "20001",
|
|
"sender_name": "历史会话",
|
|
"sender_avatar": "https://example.test/a.png",
|
|
"content": "历史消息",
|
|
"unread_count": 0,
|
|
},
|
|
{
|
|
"conversation_id": "0:1:10001:20002",
|
|
"peer_uid": "20002",
|
|
"sender_name": "未读会话",
|
|
"sender_avatar": "https://example.test/b.png",
|
|
"content": "新消息",
|
|
"unread_count": 1,
|
|
},
|
|
],
|
|
[
|
|
{
|
|
"conversation_id": "0:1:10001:20001",
|
|
"peer_uid": "20001",
|
|
"sender_name": "历史会话",
|
|
"sender_avatar": "https://example.test/a.png",
|
|
"content": "历史消息",
|
|
"unread_count": 0,
|
|
},
|
|
{
|
|
"conversation_id": "0:1:10001:20002",
|
|
"peer_uid": "20002",
|
|
"sender_name": "未读会话",
|
|
"sender_avatar": "https://example.test/b.png",
|
|
"content": "新消息",
|
|
"unread_count": 0,
|
|
},
|
|
],
|
|
[
|
|
{
|
|
"conversation_id": "0:1:10001:20001",
|
|
"peer_uid": "20001",
|
|
"sender_name": "历史会话",
|
|
"sender_avatar": "https://example.test/a.png",
|
|
"content": "真正发生变化",
|
|
"unread_count": 0,
|
|
},
|
|
{
|
|
"conversation_id": "0:1:10001:20002",
|
|
"peer_uid": "20002",
|
|
"sender_name": "未读会话",
|
|
"sender_avatar": "https://example.test/b.png",
|
|
"content": "新消息",
|
|
"unread_count": 0,
|
|
},
|
|
],
|
|
]
|
|
|
|
class _HttpClient:
|
|
def __init__(self):
|
|
self.get_conversations = AsyncMock(side_effect=snapshots)
|
|
self.conversation_list_unsupported = False
|
|
self.enter_count = 0
|
|
self.exit_count = 0
|
|
|
|
async def __aenter__(self):
|
|
self.enter_count += 1
|
|
return self
|
|
|
|
async def __aexit__(self, *_args):
|
|
self.exit_count += 1
|
|
return False
|
|
|
|
http = _HttpClient()
|
|
service = DouyinImService(
|
|
session=DouyinImSession(cookies={"sessionid": "test"}, my_uid=10001),
|
|
match_reply=AsyncMock(return_value=None),
|
|
log_fn=AsyncMock(),
|
|
account_id=9,
|
|
)
|
|
service._handle_incoming = AsyncMock()
|
|
controller = _Controller()
|
|
|
|
with (
|
|
patch.object(
|
|
service_module,
|
|
"get_traffic_controller",
|
|
return_value=controller,
|
|
),
|
|
patch.object(service_module, "DouyinImHttpClient", return_value=http) as factory,
|
|
):
|
|
# Startup indexes both previews, but only the unread conversation
|
|
# is allowed to enter the reply path.
|
|
await service._poll_conversations(initial=True)
|
|
self.assertEqual(service._handle_incoming.await_count, 1)
|
|
self.assertEqual(
|
|
service._handle_incoming.await_args.args[0]["peer_uid"],
|
|
"20002",
|
|
)
|
|
|
|
# The first normal reconciliation sees the exact same previews;
|
|
# it must not merely defer a historical-message reply explosion.
|
|
service._handle_incoming.reset_mock()
|
|
await service._poll_conversations()
|
|
service._handle_incoming.assert_not_awaited()
|
|
|
|
# A real preview transition is processed even if unread_count is
|
|
# unavailable/zero on the upstream response.
|
|
await service._poll_conversations()
|
|
service._handle_incoming.assert_awaited_once()
|
|
self.assertEqual(
|
|
service._handle_incoming.await_args.args[0]["peer_uid"],
|
|
"20001",
|
|
)
|
|
|
|
self.assertEqual(factory.call_count, 3)
|
|
self.assertEqual(http.enter_count, 3)
|
|
self.assertEqual(http.exit_count, 3)
|
|
self.assertEqual(controller.startup_flags, [True, False, False])
|
|
|
|
async def test_ready_is_not_blocked_by_slow_initial_unread_handler(self):
|
|
events: list[str] = []
|
|
ready = asyncio.Event()
|
|
handler_started = asyncio.Event()
|
|
handler_cancelled = asyncio.Event()
|
|
never_release = asyncio.Event()
|
|
|
|
def on_ready():
|
|
events.append("ready")
|
|
ready.set()
|
|
|
|
async def slow_handler(_message):
|
|
events.append("handler")
|
|
handler_started.set()
|
|
try:
|
|
await never_release.wait()
|
|
except asyncio.CancelledError:
|
|
handler_cancelled.set()
|
|
raise
|
|
|
|
class _WsClient:
|
|
connected = False
|
|
|
|
def __init__(self, *_args, **_kwargs):
|
|
self.start = AsyncMock()
|
|
self.stop = AsyncMock()
|
|
|
|
service = DouyinImService(
|
|
session=DouyinImSession(cookies={"sessionid": "test"}, my_uid=10001),
|
|
match_reply=AsyncMock(return_value=None),
|
|
log_fn=AsyncMock(),
|
|
account_id=901,
|
|
on_ready=on_ready,
|
|
)
|
|
service._verify_account_uid = AsyncMock()
|
|
service._poll_conversations = AsyncMock(
|
|
return_value=[
|
|
{
|
|
"conversation_id": "0:1:10001:29001",
|
|
"content": "startup unread",
|
|
"unread_count": 1,
|
|
}
|
|
]
|
|
)
|
|
service._handle_incoming = AsyncMock(side_effect=slow_handler)
|
|
service._reply_queue.start = AsyncMock()
|
|
service._reply_queue.stop = AsyncMock()
|
|
|
|
with (
|
|
patch.object(service_module, "DouyinImWsClient", _WsClient),
|
|
patch.object(service_module, "ensure_frontier_ws"),
|
|
patch.object(service_module.system_logger, "record"),
|
|
patch(
|
|
"rpa_engine.douyin_im.emoji_pack.is_fresh",
|
|
return_value=True,
|
|
),
|
|
):
|
|
run_task = asyncio.create_task(service.run())
|
|
try:
|
|
# If initial unread were still processed inline, this wait
|
|
# would time out because slow_handler never completes.
|
|
await asyncio.wait_for(ready.wait(), timeout=0.3)
|
|
await asyncio.wait_for(handler_started.wait(), timeout=0.3)
|
|
self.assertEqual(events[:2], ["ready", "handler"])
|
|
service._poll_conversations.assert_awaited_once_with(
|
|
initial=True,
|
|
defer_handlers=True,
|
|
)
|
|
|
|
# Account stop cancels its active deferred handler instead of
|
|
# leaving work detached from the service lifecycle.
|
|
await asyncio.wait_for(service.stop(), timeout=0.5)
|
|
self.assertTrue(handler_cancelled.is_set())
|
|
finally:
|
|
run_task.cancel()
|
|
await asyncio.gather(run_task, return_exceptions=True)
|
|
await service_module._shutdown_initial_unread_dispatcher()
|
|
|
|
async def test_initial_unread_dispatcher_has_process_wide_concurrency_limit(self):
|
|
dispatcher = service_module._InitialUnreadDispatcher(concurrency=2)
|
|
active = 0
|
|
maximum_active = 0
|
|
processed: list[int] = []
|
|
two_started = asyncio.Event()
|
|
release = asyncio.Event()
|
|
|
|
def make_service(account_id: int):
|
|
async def handle(_message):
|
|
nonlocal active, maximum_active
|
|
active += 1
|
|
maximum_active = max(maximum_active, active)
|
|
if active == 2:
|
|
two_started.set()
|
|
try:
|
|
await release.wait()
|
|
processed.append(account_id)
|
|
finally:
|
|
active -= 1
|
|
|
|
return SimpleNamespace(
|
|
account_id=account_id,
|
|
_running=True,
|
|
_handle_incoming=handle,
|
|
)
|
|
|
|
services = [make_service(index + 1) for index in range(6)]
|
|
try:
|
|
for service in services:
|
|
await dispatcher.submit(service, [{"unread_count": 1}])
|
|
|
|
await asyncio.wait_for(two_started.wait(), timeout=0.3)
|
|
await asyncio.sleep(0.03)
|
|
self.assertEqual(maximum_active, 2)
|
|
self.assertEqual(processed, [])
|
|
|
|
release.set()
|
|
await asyncio.wait_for(dispatcher.join(), timeout=0.5)
|
|
self.assertEqual(maximum_active, 2)
|
|
self.assertCountEqual(processed, range(1, 7))
|
|
finally:
|
|
await dispatcher.stop()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|