591 lines
21 KiB
Python
591 lines
21 KiB
Python
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_immediate_if_idle_runs_first_now_then_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()
|
|
first_started = asyncio.Event()
|
|
release_first = asyncio.Event()
|
|
|
|
def callback_for(index: int):
|
|
async def callback() -> None:
|
|
calls.append((index, loop.time() - started_at))
|
|
if index == 0:
|
|
first_started.set()
|
|
await release_first.wait()
|
|
if len(calls) == 3:
|
|
finished.set()
|
|
|
|
return callback
|
|
|
|
await queue.enqueue(
|
|
interval,
|
|
callback_for(0),
|
|
description="0",
|
|
immediate_if_idle=True,
|
|
)
|
|
await asyncio.wait_for(first_started.wait(), timeout=0.1)
|
|
for index in (1, 2):
|
|
await queue.enqueue(
|
|
interval,
|
|
callback_for(index),
|
|
description=str(index),
|
|
immediate_if_idle=True,
|
|
)
|
|
release_first.set()
|
|
|
|
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, (0, interval, interval * 2)):
|
|
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_immediate_if_idle_does_not_bypass_active_send(self):
|
|
queue = await self._start_queue(account_id=102)
|
|
interval = 1.0
|
|
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(
|
|
interval,
|
|
first_callback,
|
|
description="first",
|
|
immediate_if_idle=True,
|
|
)
|
|
await asyncio.wait_for(first_started.wait(), timeout=0.1)
|
|
await queue.enqueue(
|
|
interval,
|
|
second_callback,
|
|
description="second",
|
|
immediate_if_idle=True,
|
|
)
|
|
|
|
snapshot = await queue.snapshot()
|
|
self.assertEqual([item["status"] for item in snapshot], ["sending", "waiting"])
|
|
self.assertEqual(snapshot[0]["interval_seconds"], 0)
|
|
self.assertEqual(snapshot[1]["interval_seconds"], int(interval))
|
|
|
|
release_first.set()
|
|
await asyncio.sleep(0.02)
|
|
self.assertFalse(second_started.is_set())
|
|
|
|
async def test_immediate_if_idle_resets_after_queue_drains(self):
|
|
queue = await self._start_queue(account_id=103)
|
|
interval = 1.0
|
|
first_finished = asyncio.Event()
|
|
second_started = asyncio.Event()
|
|
release_second = asyncio.Event()
|
|
|
|
async def first_callback() -> None:
|
|
first_finished.set()
|
|
|
|
async def second_callback() -> None:
|
|
second_started.set()
|
|
await release_second.wait()
|
|
|
|
await queue.enqueue(
|
|
interval,
|
|
first_callback,
|
|
description="first wave",
|
|
immediate_if_idle=True,
|
|
)
|
|
await asyncio.wait_for(first_finished.wait(), timeout=0.1)
|
|
await self._wait_until_idle(queue)
|
|
|
|
await queue.enqueue(
|
|
interval,
|
|
second_callback,
|
|
description="second wave",
|
|
immediate_if_idle=True,
|
|
)
|
|
await asyncio.wait_for(second_started.wait(), timeout=0.1)
|
|
snapshot = await queue.snapshot()
|
|
self.assertEqual(len(snapshot), 1)
|
|
self.assertEqual(snapshot[0]["status"], "sending")
|
|
self.assertEqual(snapshot[0]["interval_seconds"], 0)
|
|
release_second.set()
|
|
|
|
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_on_zero_slot_does_not_claim_later_jobs_shifted(self):
|
|
queue = await self._start_queue(account_id=512)
|
|
interval = 1.0
|
|
|
|
async def noop() -> None:
|
|
pass
|
|
|
|
# Both enqueues complete without yielding to the consumer, preserving
|
|
# the narrow management-API window where the zero-slot first job is
|
|
# still waiting and can be selected by send-now.
|
|
await queue.enqueue(
|
|
interval,
|
|
noop,
|
|
description="immediate first",
|
|
immediate_if_idle=True,
|
|
)
|
|
await queue.enqueue(
|
|
interval,
|
|
noop,
|
|
description="scheduled second",
|
|
immediate_if_idle=True,
|
|
)
|
|
before = await queue.snapshot()
|
|
second_before = next(
|
|
item for item in before if item["description"] == "scheduled second"
|
|
)
|
|
|
|
result = await queue.send_now(before[0]["job_id"])
|
|
after = await queue.snapshot()
|
|
second_after = next(
|
|
item for item in after if item["description"] == "scheduled second"
|
|
)
|
|
|
|
self.assertEqual(result["status"], "accepted")
|
|
self.assertEqual(result["shifted_count"], 0)
|
|
second_due_before = datetime.fromisoformat(
|
|
second_before["scheduled_at"]
|
|
).timestamp()
|
|
second_due_after = datetime.fromisoformat(second_after["scheduled_at"]).timestamp()
|
|
self.assertAlmostEqual(second_due_after, second_due_before, delta=0.01)
|
|
|
|
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()
|