更新

backend/tests/test_conversation_poll_bandwidth.py

@@ -1,10 +1,12 @@
 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
 
 
@@ -109,6 +111,42 @@ class ConversationPollBandwidthTests(unittest.IsolatedAsyncioTestCase):
         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
@@ -141,7 +179,260 @@ class ConversationPollBandwidthTests(unittest.IsolatedAsyncioTestCase):
             await service._poll_conversations()
 
         http.get_conversations.assert_awaited_once_with(enrich_profiles=False)
-        service._index_conversations.assert_awaited_once_with([])
+        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.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__":