import asyncio import sys import unittest from datetime import datetime from pathlib import Path BACKEND_DIR = Path(__file__).resolve().parents[1] if str(BACKEND_DIR) not in sys.path: sys.path.insert(0, str(BACKEND_DIR)) from rpa_engine.douyin_im.reply_queue import AccountReplyQueue class AccountReplyQueueTests(unittest.IsolatedAsyncioTestCase): async def _start_queue(self, account_id: int, **kwargs) -> AccountReplyQueue: queue = AccountReplyQueue(account_id=account_id, **kwargs) await queue.start() self.addAsyncCleanup(queue.stop) return queue async def _wait_until_idle( self, queue: AccountReplyQueue, timeout: float = 0.5, ) -> None: loop = asyncio.get_running_loop() deadline = loop.time() + timeout while queue.pending_count and loop.time() < deadline: await asyncio.sleep(0.002) self.assertEqual(queue.pending_count, 0) async def test_one_account_runs_three_jobs_at_successive_fifo_slots(self): queue = await self._start_queue(account_id=101) interval = 0.05 loop = asyncio.get_running_loop() started_at = loop.time() calls: list[tuple[int, float]] = [] finished = asyncio.Event() def callback_for(index: int): async def callback() -> None: calls.append((index, loop.time() - started_at)) if len(calls) == 3: finished.set() return callback for index in range(3): await queue.enqueue(interval, callback_for(index), description=str(index)) await asyncio.wait_for(finished.wait(), timeout=0.75) await self._wait_until_idle(queue) self.assertEqual([index for index, _ in calls], [0, 1, 2]) elapsed = [timestamp for _, timestamp in calls] for timestamp, expected in zip(elapsed, (interval, interval * 2, interval * 3)): self.assertGreaterEqual(timestamp, expected - 0.015) self.assertLess(timestamp, expected + 0.15) self.assertGreaterEqual(elapsed[1] - elapsed[0], interval - 0.02) self.assertGreaterEqual(elapsed[2] - elapsed[1], interval - 0.02) async def test_separate_account_queues_reach_first_slot_without_blocking(self): first_queue = await self._start_queue(account_id=201) second_queue = await self._start_queue(account_id=202) interval = 0.04 loop = asyncio.get_running_loop() started_at = loop.time() first_started = asyncio.Event() release_first = asyncio.Event() second_finished = asyncio.Event() timestamps: dict[str, float] = {} async def first_callback() -> None: timestamps["first"] = loop.time() - started_at first_started.set() await release_first.wait() async def second_callback() -> None: timestamps["second"] = loop.time() - started_at second_finished.set() await asyncio.gather( first_queue.enqueue(interval, first_callback, description="first account"), second_queue.enqueue(interval, second_callback, description="second account"), ) await asyncio.wait_for(first_started.wait(), timeout=0.4) await asyncio.wait_for(second_finished.wait(), timeout=0.4) self.assertFalse(release_first.is_set()) self.assertLess(abs(timestamps["first"] - timestamps["second"]), 0.05) release_first.set() await self._wait_until_idle(first_queue) await self._wait_until_idle(second_queue) async def test_stop_before_due_discards_pending_job(self): queue = await self._start_queue(account_id=301) callback_called = asyncio.Event() async def callback() -> None: callback_called.set() await queue.enqueue(0.12, callback, description="must be discarded") await asyncio.sleep(0.02) self.assertEqual(queue.pending_count, 1) await queue.stop() self.assertEqual(queue.pending_count, 0) await asyncio.sleep(0.13) self.assertFalse(callback_called.is_set()) async def test_callback_failure_does_not_block_following_job(self): errors: list[tuple[str, str]] = [] def on_error(description: str, exc: BaseException) -> None: errors.append((description, type(exc).__name__)) queue = await self._start_queue(account_id=401, on_error=on_error) calls: list[str] = [] second_finished = asyncio.Event() async def failing_callback() -> None: calls.append("first") raise ValueError("expected test failure") async def following_callback() -> None: calls.append("second") second_finished.set() with self.assertLogs("douyin_im.reply_queue", level="ERROR"): await queue.enqueue(0.03, failing_callback, description="first") await queue.enqueue(0.03, following_callback, description="second") await asyncio.wait_for(second_finished.wait(), timeout=0.5) await self._wait_until_idle(queue) self.assertEqual(calls, ["first", "second"]) self.assertEqual(errors, [("first", "ValueError")]) async def test_snapshot_exposes_waiting_job_details(self): queue = await self._start_queue(account_id=501) async def callback() -> None: pass await queue.enqueue( 0.2, callback, description="回复 测试用户", details={ "sender_name": "测试用户", "conversation_id": "conv-501", "incoming_content": "你好", "replies": ["您好"], }, ) items = await queue.snapshot() self.assertEqual(len(items), 1) self.assertEqual(items[0]["position"], 1) self.assertEqual(items[0]["status"], "waiting") self.assertEqual(items[0]["sender_name"], "测试用户") self.assertEqual(items[0]["incoming_content"], "你好") self.assertEqual(items[0]["replies"], ["您好"]) self.assertNotIn("callback", items[0]) async def test_send_now_wakes_first_job_and_moves_later_slots_forward(self): queue = await self._start_queue(account_id=502) interval = 0.2 sent = asyncio.Event() async def first_callback() -> None: sent.set() async def noop() -> None: pass await queue.enqueue(interval, first_callback, description="first") await queue.enqueue(interval, noop, description="second") await queue.enqueue(interval, noop, description="third") before = await queue.snapshot() result = await queue.send_now(before[0]["job_id"]) self.assertEqual(result["status"], "accepted") self.assertEqual(result["shifted_count"], 2) await asyncio.wait_for(sent.wait(), timeout=0.1) await asyncio.sleep(0) after = await queue.snapshot() self.assertEqual([item["description"] for item in after], ["second", "third"]) for old, new in zip(before[1:], after): old_due = datetime.fromisoformat(old["scheduled_at"]).timestamp() new_due = datetime.fromisoformat(new["scheduled_at"]).timestamp() self.assertAlmostEqual(old_due - new_due, interval, delta=0.05) async def test_send_now_middle_runs_first_and_only_shifts_jobs_behind_it(self): queue = await self._start_queue(account_id=503) interval = 0.12 calls: list[str] = [] all_sent = asyncio.Event() def callback_for(name: str): async def callback() -> None: calls.append(name) if len(calls) == 3: all_sent.set() return callback for name in ("first", "middle", "last"): await queue.enqueue(interval, callback_for(name), description=name) before = await queue.snapshot() result = await queue.send_now(before[1]["job_id"]) self.assertEqual(result["shifted_count"], 1) await asyncio.wait_for(all_sent.wait(), timeout=0.6) self.assertEqual(calls, ["middle", "first", "last"]) first_before = datetime.fromisoformat(before[0]["scheduled_at"]).timestamp() # The first normal job keeps its original slot; this is also validated by execution order. self.assertGreater(first_before, datetime.now().timestamp() - 1) async def test_send_now_does_not_overlap_active_send(self): queue = await self._start_queue(account_id=504) first_started = asyncio.Event() release_first = asyncio.Event() second_started = asyncio.Event() async def first_callback() -> None: first_started.set() await release_first.wait() async def second_callback() -> None: second_started.set() await queue.enqueue(0.01, first_callback, description="first") await queue.enqueue(0.2, second_callback, description="second") await asyncio.wait_for(first_started.wait(), timeout=0.2) items = await queue.snapshot() second = next(item for item in items if item["description"] == "second") result = await queue.send_now(second["job_id"]) self.assertEqual(result["status"], "accepted") await asyncio.sleep(0.03) self.assertFalse(second_started.is_set()) release_first.set() await asyncio.wait_for(second_started.wait(), timeout=0.2) async def test_send_now_tail_releases_slot_for_next_enqueue(self): queue = await self._start_queue(account_id=505) interval = 0.2 async def noop() -> None: pass for name in ("first", "second", "tail"): await queue.enqueue(interval, noop, description=name) before = await queue.snapshot() old_tail_due = datetime.fromisoformat(before[2]["scheduled_at"]).timestamp() await queue.send_now(before[2]["job_id"]) await queue.enqueue(interval, noop, description="new-tail") after = await queue.snapshot() new_tail = next(item for item in after if item["description"] == "new-tail") new_tail_due = datetime.fromisoformat(new_tail["scheduled_at"]).timestamp() self.assertAlmostEqual(new_tail_due, old_tail_due, delta=0.05) async def test_repeated_send_now_is_idempotent(self): queue = await self._start_queue(account_id=506) calls = 0 finished = asyncio.Event() async def callback() -> None: nonlocal calls calls += 1 finished.set() await queue.enqueue(0.2, callback, description="once") item = (await queue.snapshot())[0] first = await queue.send_now(item["job_id"]) second = await queue.send_now(item["job_id"]) self.assertEqual(first["status"], "accepted") self.assertIn(second["status"], {"already_requested", "already_sending"}) await asyncio.wait_for(finished.wait(), timeout=0.1) await asyncio.sleep(0.02) self.assertEqual(calls, 1) async def test_merge_pending_keeps_one_job_slot_and_callback(self): queue = await self._start_queue(account_id=507) calls = 0 finished = asyncio.Event() async def callback() -> None: nonlocal calls calls += 1 finished.set() await queue.enqueue( 0.3, callback, description="conversation", details={ "incoming_content": "first", "incoming_contents": ["first"], "message_count": 1, "replies": ["one reply"], }, merge_key="conv:507", ) before = (await queue.snapshot())[0] def append_second(details): details["incoming_content"] = "second" details["incoming_contents"].append("second") details["message_count"] = 2 return details result = await queue.merge_pending("conv:507", append_second) after = await queue.snapshot() self.assertEqual(result["status"], "merged") self.assertEqual(result["position"], 1) self.assertEqual(len(after), 1) self.assertEqual(after[0]["job_id"], before["job_id"]) self.assertEqual(after[0]["incoming_contents"], ["first", "second"]) self.assertEqual(after[0]["replies"], ["one reply"]) old_due = datetime.fromisoformat(before["scheduled_at"]).timestamp() new_due = datetime.fromisoformat(after[0]["scheduled_at"]).timestamp() self.assertAlmostEqual(old_due, new_due, delta=0.02) # Management snapshots must not expose the queue's mutable nested list. after[0]["incoming_contents"].append("external mutation") self.assertEqual( (await queue.snapshot())[0]["incoming_contents"], ["first", "second"], ) await queue.send_now(before["job_id"]) await asyncio.wait_for(finished.wait(), timeout=0.2) self.assertEqual(calls, 1) async def test_urgent_job_can_still_merge_before_sending(self): queue = await self._start_queue(account_id=508) blocker_started = asyncio.Event() release_blocker = asyncio.Event() merged_sent = asyncio.Event() merged_calls = 0 async def blocker() -> None: blocker_started.set() await release_blocker.wait() async def merged_callback() -> None: nonlocal merged_calls merged_calls += 1 merged_sent.set() await queue.enqueue(0.01, blocker, description="blocker") await queue.enqueue( 0.3, merged_callback, description="mergeable", details={"incoming_content": "first", "incoming_contents": ["first"]}, merge_key="conv:508", ) await asyncio.wait_for(blocker_started.wait(), timeout=0.2) mergeable = next( item for item in await queue.snapshot() if item["description"] == "mergeable" ) await queue.send_now(mergeable["job_id"]) def append_second(details): details["incoming_contents"].append("second") details["message_count"] = 2 return details result = await queue.merge_pending("conv:508", append_second) self.assertEqual(result["status"], "merged") self.assertEqual(result["queue_status"], "ready") release_blocker.set() await asyncio.wait_for(merged_sent.wait(), timeout=0.2) self.assertEqual(merged_calls, 1) async def test_active_job_is_sealed_against_merge(self): queue = await self._start_queue(account_id=509) started = asyncio.Event() release = asyncio.Event() async def callback() -> None: started.set() await release.wait() await queue.enqueue( 0.01, callback, details={"incoming_content": "first", "incoming_contents": ["first"]}, merge_key="conv:509", ) await asyncio.wait_for(started.wait(), timeout=0.2) result = await queue.merge_pending("conv:509", lambda details: details) self.assertEqual(result["status"], "not_found") release.set() async def test_merge_keys_bridge_partial_peer_and_conversation_ids(self): queue = await self._start_queue(account_id=510) async def callback() -> None: pass await queue.enqueue( 0.3, callback, details={"incoming_contents": ["first"], "message_count": 1}, merge_keys=("peer:123",), ) def append_second(details): details["incoming_contents"].append("second") details["message_count"] = 2 return details result = await queue.merge_pending( ("conv:510", "peer:123"), append_second, ) self.assertEqual(result["status"], "merged") self.assertEqual(queue.pending_count, 1) self.assertEqual( (await queue.snapshot())[0]["incoming_contents"], ["first", "second"], ) async def test_different_conversation_ids_do_not_merge_on_shared_peer(self): queue = await self._start_queue(account_id=511) async def callback() -> None: pass await queue.enqueue( 0.3, callback, details={"incoming_contents": ["first"]}, merge_keys=("conv:first", "peer:123"), ) result = await queue.merge_pending( ("conv:second", "peer:123"), lambda details: details, ) self.assertEqual(result["status"], "not_found") self.assertEqual(queue.pending_count, 1) if __name__ == "__main__": unittest.main()