Files
agentscope-ai--agentscope/tests/in_memory_message_bus_test.py
wehub-resource-sync c3bf08ac8d
K8s Workspace Integration Tests / k8s-workspace-tests (push) Waiting to run
Pre-commit / run (ubuntu-latest) (push) Waiting to run
Python Unittest Coverage / test (macos-15, 3.11) (push) Waiting to run
Python Unittest Coverage / test (ubuntu-latest, 3.11) (push) Waiting to run
Python Unittest Coverage / test (windows-latest, 3.11) (push) Waiting to run
Web UI / check (push) Waiting to run
chore: import upstream snapshot with attribution
2026-07-13 12:39:27 +08:00

504 lines
19 KiB
Python

# -*- coding: utf-8 -*-
"""Tests for :class:`InMemoryMessageBus`.
The same abstract surface exercised in ``service_message_bus_test.py``
(queue / log / pubsub / lock / registry) is tested here against the
pure-Python in-memory backend, plus the domain helpers inherited from
the base :class:`MessageBus` class.
No external dependencies (no Redis, no fakeredis) — just asyncio.
"""
import asyncio
from contextlib import AsyncExitStack
from unittest import IsolatedAsyncioTestCase
from agentscope.app.message_bus import InMemoryMessageBus
class TestQueuePrimitive(IsolatedAsyncioTestCase):
"""Mode A — ``queue_push`` + ``queue_drain`` semantics."""
async def asyncSetUp(self) -> None:
self._stack = AsyncExitStack()
self.bus = await self._stack.enter_async_context(
InMemoryMessageBus(),
)
async def asyncTearDown(self) -> None:
await self._stack.aclose()
async def test_push_drain_returns_payloads_in_order(self) -> None:
"""Entries pushed in order come back out in order, once each."""
await self.bus.queue_push("k", {"i": 1})
await self.bus.queue_push("k", {"i": 2})
entries = await self.bus.queue_drain("k", max_count=10)
self.assertEqual([p for _id, p in entries], [{"i": 1}, {"i": 2}])
async def test_drain_is_destructive(self) -> None:
"""A drained entry is gone; a second drain yields nothing."""
await self.bus.queue_push("k", {"x": 1})
await self.bus.queue_drain("k", max_count=10)
self.assertEqual(await self.bus.queue_drain("k", max_count=10), [])
async def test_drain_respects_max_count(self) -> None:
"""``max_count`` caps the batch size; remaining entries persist."""
for i in range(5):
await self.bus.queue_push("k", {"i": i})
first = await self.bus.queue_drain("k", max_count=3)
rest = await self.bus.queue_drain("k", max_count=10)
self.assertEqual([p["i"] for _id, p in first], [0, 1, 2])
self.assertEqual([p["i"] for _id, p in rest], [3, 4])
async def test_drain_empty_queue_returns_empty(self) -> None:
"""Draining a key that was never pushed returns an empty list."""
self.assertEqual(await self.bus.queue_drain("nope"), [])
async def test_push_returns_unique_ids(self) -> None:
"""Each ``queue_push`` returns a distinct entry id."""
id1 = await self.bus.queue_push("k", {"a": 1})
id2 = await self.bus.queue_push("k", {"a": 2})
self.assertNotEqual(id1, id2)
async def test_queue_delete_removes_all(self) -> None:
"""``queue_delete`` drops the entire queue."""
await self.bus.queue_push("k", {"i": 1})
await self.bus.queue_push("k", {"i": 2})
await self.bus.queue_delete("k")
self.assertEqual(await self.bus.queue_drain("k", max_count=10), [])
async def test_queue_delete_missing_is_noop(self) -> None:
"""Deleting a non-existent queue does not raise."""
await self.bus.queue_delete("never-existed")
async def test_queue_isolation_between_keys(self) -> None:
"""Pushes to different keys are independent."""
await self.bus.queue_push("a", {"x": 1})
await self.bus.queue_push("b", {"x": 2})
a = await self.bus.queue_drain("a", max_count=10)
b = await self.bus.queue_drain("b", max_count=10)
self.assertEqual([p for _id, p in a], [{"x": 1}])
self.assertEqual([p for _id, p in b], [{"x": 2}])
class TestLogPrimitive(IsolatedAsyncioTestCase):
"""Mode C — replay log: append / read with cursor / trim."""
async def asyncSetUp(self) -> None:
self._stack = AsyncExitStack()
self.bus = await self._stack.enter_async_context(
InMemoryMessageBus(),
)
async def asyncTearDown(self) -> None:
await self._stack.aclose()
async def test_read_returns_everything_when_no_cursor(self) -> None:
"""Without a ``since`` cursor, the whole log comes back."""
await self.bus.log_append("k", {"i": 1})
await self.bus.log_append("k", {"i": 2})
entries = await self.bus.log_read("k")
self.assertEqual([p["i"] for _id, p in entries], [1, 2])
async def test_read_with_cursor_is_exclusive(self) -> None:
"""``since=last_id`` skips that id and returns only newer."""
await self.bus.log_append("k", {"i": 1})
await self.bus.log_append("k", {"i": 2})
await self.bus.log_append("k", {"i": 3})
all_entries = await self.bus.log_read("k")
cursor = all_entries[1][0] # id of entry 2
rest = await self.bus.log_read("k", since=cursor)
self.assertEqual([p["i"] for _id, p in rest], [3])
async def test_read_respects_max_count(self) -> None:
"""``max_count`` caps the batch; remaining entries persist."""
for i in range(5):
await self.bus.log_append("k", {"i": i})
first = await self.bus.log_read("k", max_count=3)
self.assertEqual([p["i"] for _id, p in first], [0, 1, 2])
async def test_read_is_non_destructive(self) -> None:
"""Multiple reads on the same log return the same data."""
await self.bus.log_append("k", {"i": 1})
r1 = await self.bus.log_read("k")
r2 = await self.bus.log_read("k")
self.assertEqual(
[p["i"] for _id, p in r1],
[p["i"] for _id, p in r2],
)
async def test_read_empty_log(self) -> None:
"""Reading a log that never had entries returns ``[]``."""
self.assertEqual(await self.bus.log_read("nope"), [])
async def test_read_all_before_cursor(self) -> None:
"""When all entries are at or before the cursor, result is ``[]``."""
id2 = await self.bus.log_append("k", {"i": 2})
self.assertEqual(await self.bus.log_read("k", since=id2), [])
async def test_trim_without_before_drops_entire_log(self) -> None:
"""``log_trim(key)`` empties the log."""
await self.bus.log_append("k", {"i": 1})
await self.bus.log_append("k", {"i": 2})
await self.bus.log_trim("k")
self.assertEqual(await self.bus.log_read("k"), [])
async def test_trim_with_before_id_keeps_newer(self) -> None:
"""``log_trim(key, before_id)`` drops older entries only."""
await self.bus.log_append("k", {"i": 1})
id2 = await self.bus.log_append("k", {"i": 2})
await self.bus.log_append("k", {"i": 3})
await self.bus.log_trim("k", before_id=id2)
entries = await self.bus.log_read("k")
self.assertEqual([p["i"] for _id, p in entries], [2, 3])
async def test_trim_missing_key_is_noop(self) -> None:
"""Trimming a non-existent log does not raise."""
await self.bus.log_trim("nope")
async def test_max_len_caps_log_size(self) -> None:
"""``max_len`` on ``log_append`` trims older entries when the
log exceeds the cap."""
for i in range(10):
await self.bus.log_append("k", {"i": i}, max_len=5)
entries = await self.bus.log_read("k", max_count=100)
self.assertLessEqual(len(entries), 5)
# The newest entries must survive.
self.assertEqual(entries[-1][1]["i"], 9)
class TestPubSubPrimitive(IsolatedAsyncioTestCase):
"""Mode D — transient broadcast: publish / subscribe."""
async def asyncSetUp(self) -> None:
self._stack = AsyncExitStack()
self.bus = await self._stack.enter_async_context(
InMemoryMessageBus(),
)
async def asyncTearDown(self) -> None:
await self._stack.aclose()
async def test_subscribe_receives_messages_published_after_ready(
self,
) -> None:
"""Subscribers receive payloads published after the subscription
is established."""
ready = asyncio.Event()
received: list[dict] = []
async def _consumer() -> None:
async for payload in self.bus.subscribe(
"ch",
on_ready=ready.set,
):
received.append(payload)
if len(received) == 2:
break
task = asyncio.create_task(_consumer())
await asyncio.wait_for(ready.wait(), timeout=2.0)
await self.bus.publish("ch", {"i": 1})
await self.bus.publish("ch", {"i": 2})
await asyncio.wait_for(task, timeout=2.0)
self.assertEqual([p["i"] for p in received], [1, 2])
async def test_publish_without_subscribers_is_noop(self) -> None:
"""Publishing when no one is listening does not raise."""
await self.bus.publish("ch", {"i": 1})
async def test_multiple_subscribers_each_receive(self) -> None:
"""All active subscribers on a channel receive the payload."""
ready1 = asyncio.Event()
ready2 = asyncio.Event()
r1: list[dict] = []
r2: list[dict] = []
async def _c1() -> None:
async for payload in self.bus.subscribe(
"ch",
on_ready=ready1.set,
):
r1.append(payload)
break
async def _c2() -> None:
async for payload in self.bus.subscribe(
"ch",
on_ready=ready2.set,
):
r2.append(payload)
break
t1 = asyncio.create_task(_c1())
t2 = asyncio.create_task(_c2())
await asyncio.wait_for(ready1.wait(), timeout=2.0)
await asyncio.wait_for(ready2.wait(), timeout=2.0)
await self.bus.publish("ch", {"x": 42})
await asyncio.wait_for(t1, timeout=2.0)
await asyncio.wait_for(t2, timeout=2.0)
self.assertEqual(r1, [{"x": 42}])
self.assertEqual(r2, [{"x": 42}])
class TestLockPrimitive(IsolatedAsyncioTestCase):
"""Mode E — distributed mutex (process-local asyncio.Lock)."""
async def asyncSetUp(self) -> None:
self._stack = AsyncExitStack()
self.bus = await self._stack.enter_async_context(
InMemoryMessageBus(),
)
async def asyncTearDown(self) -> None:
await self._stack.aclose()
async def test_is_locked_reflects_acquire_release(self) -> None:
"""``is_locked`` flips to True while the body runs and back to
False once the context exits."""
self.assertFalse(await self.bus.is_locked("k"))
async with self.bus.acquire_lock("k", ttl_secs=10):
self.assertTrue(await self.bus.is_locked("k"))
self.assertFalse(await self.bus.is_locked("k"))
async def test_second_acquirer_waits_until_release(self) -> None:
"""A second ``acquire_lock`` on the same key blocks until the
first releases."""
order: list[str] = []
async def _holder() -> None:
async with self.bus.acquire_lock("k", ttl_secs=10):
order.append("first-in")
await asyncio.sleep(0.05)
order.append("first-out")
async def _challenger() -> None:
await asyncio.sleep(0.005)
async with self.bus.acquire_lock("k", ttl_secs=10):
order.append("second-in")
await asyncio.gather(_holder(), _challenger())
self.assertEqual(
order,
["first-in", "first-out", "second-in"],
)
class TestRegistryPrimitive(IsolatedAsyncioTestCase):
"""Mode F — ``registry_*`` hash-keyed namespace operations."""
async def asyncSetUp(self) -> None:
self._stack = AsyncExitStack()
self.bus = await self._stack.enter_async_context(
InMemoryMessageBus(),
)
async def asyncTearDown(self) -> None:
await self._stack.aclose()
async def test_set_then_exists_and_getall(self) -> None:
"""``registry_set`` stores a field; ``exists`` and ``getall``
round-trip correctly."""
await self.bus.registry_set("ns", "f1", "v1")
await self.bus.registry_set("ns", "f2", "v2")
self.assertTrue(await self.bus.registry_exists("ns", "f1"))
self.assertTrue(await self.bus.registry_exists("ns", "f2"))
self.assertFalse(await self.bus.registry_exists("ns", "missing"))
self.assertFalse(await self.bus.registry_exists("other-ns", "f1"))
self.assertEqual(
await self.bus.registry_getall("ns"),
{"f1": "v1", "f2": "v2"},
)
async def test_set_overwrites_existing_field(self) -> None:
"""A second ``registry_set`` for the same field overwrites."""
await self.bus.registry_set("ns", "f", "v1")
await self.bus.registry_set("ns", "f", "v2")
self.assertEqual(
await self.bus.registry_getall("ns"),
{"f": "v2"},
)
async def test_del_removes_only_the_named_field(self) -> None:
"""``registry_del`` removes a single field; siblings survive."""
await self.bus.registry_set("ns", "keep", "k")
await self.bus.registry_set("ns", "drop", "d")
await self.bus.registry_del("ns", "drop")
self.assertFalse(await self.bus.registry_exists("ns", "drop"))
self.assertTrue(await self.bus.registry_exists("ns", "keep"))
async def test_del_missing_field_is_noop(self) -> None:
"""Deleting a non-existent field does not raise."""
await self.bus.registry_del("ns", "nope")
async def test_getall_on_missing_namespace_returns_empty(self) -> None:
"""``registry_getall`` for an unknown namespace returns ``{}``."""
self.assertEqual(await self.bus.registry_getall("ghost"), {})
async def test_drop_deletes_entire_namespace(self) -> None:
"""``registry_drop`` removes every field under the namespace."""
await self.bus.registry_set("ns", "f1", "v1")
await self.bus.registry_set("ns", "f2", "v2")
await self.bus.registry_drop("ns")
self.assertEqual(await self.bus.registry_getall("ns"), {})
async def test_drop_missing_namespace_is_noop(self) -> None:
"""Dropping a namespace that was never written does not raise."""
await self.bus.registry_drop("never-existed")
async def test_getall_returns_copy(self) -> None:
"""Mutating the returned dict does not affect bus state."""
await self.bus.registry_set("ns", "f", "v")
out = await self.bus.registry_getall("ns")
out["injected"] = "evil"
self.assertEqual(
await self.bus.registry_getall("ns"),
{"f": "v"},
)
class TestDomainHelpers(IsolatedAsyncioTestCase):
"""Domain helpers inherited from ``MessageBus`` work end-to-end
on the in-memory backend."""
async def asyncSetUp(self) -> None:
self._stack = AsyncExitStack()
self.bus = await self._stack.enter_async_context(
InMemoryMessageBus(),
)
async def asyncTearDown(self) -> None:
await self._stack.aclose()
async def test_session_run_trims_log_on_exit(self) -> None:
"""``session_run`` + ``session_publish_event`` + auto trim."""
sid = "s-trim"
async with self.bus.session_run(sid):
await self.bus.session_publish_event(sid, {"i": 1})
await self.bus.session_publish_event(sid, {"i": 2})
mid = await self.bus.session_read_events(sid)
self.assertEqual([p["i"] for _id, p in mid], [1, 2])
self.assertEqual(await self.bus.session_read_events(sid), [])
async def test_session_is_running_reflects_lock(self) -> None:
"""``session_is_running`` is True while inside ``session_run``."""
sid = "s-isrun"
self.assertFalse(await self.bus.session_is_running(sid))
async with self.bus.session_run(sid):
self.assertTrue(await self.bus.session_is_running(sid))
self.assertFalse(await self.bus.session_is_running(sid))
async def test_inbox_round_trip(self) -> None:
"""``inbox_push`` + ``inbox_drain`` FIFO semantics."""
sid = "s-inbox"
await self.bus.inbox_push(sid, {"hint": "a"})
await self.bus.inbox_push(sid, {"hint": "b"})
entries = await self.bus.inbox_drain(sid, max_count=10)
self.assertEqual(
[p["hint"] for _id, p in entries],
["a", "b"],
)
self.assertEqual(
await self.bus.inbox_drain(sid, max_count=10),
[],
)
async def test_enqueue_wakeup_round_trip(self) -> None:
"""``enqueue_wakeup`` → ``dequeue_wakeups`` round-trip."""
ready = asyncio.Event()
received: list[dict] = []
async def _signal_consumer() -> None:
async for payload in self.bus.subscribe_wakeup_signal(
on_ready=ready.set,
):
received.append(payload)
break
task = asyncio.create_task(_signal_consumer())
await asyncio.wait_for(ready.wait(), timeout=2.0)
await self.bus.enqueue_wakeup(
user_id="u",
session_id="s",
agent_id="a",
)
await asyncio.wait_for(task, timeout=2.0)
self.assertEqual(len(received), 1)
entries = await self.bus.dequeue_wakeups(max_count=10)
self.assertEqual(len(entries), 1)
self.assertEqual(
entries[0],
{
"user_id": "u",
"session_id": "s",
"agent_id": "a",
"kind": "wake",
"input": None,
},
)
async def test_bg_task_round_trip(self) -> None:
"""``bg_task_register / exists / list / unregister / purge``
work on the in-memory backend."""
sid = "s-bg"
self.assertFalse(await self.bus.bg_task_exists(sid, "t1"))
await self.bus.bg_task_register(sid, "t1", '{"tool":"a"}')
await self.bus.bg_task_register(sid, "t2", '{"tool":"b"}')
self.assertTrue(await self.bus.bg_task_exists(sid, "t1"))
self.assertEqual(
await self.bus.bg_task_list(sid),
{"t1": '{"tool":"a"}', "t2": '{"tool":"b"}'},
)
await self.bus.bg_task_unregister(sid, "t1")
self.assertFalse(await self.bus.bg_task_exists(sid, "t1"))
await self.bus.bg_task_purge(sid)
self.assertEqual(await self.bus.bg_task_list(sid), {})
async def test_session_purge_clears_all_bus_state(self) -> None:
"""``session_purge`` deletes the session's events + inbox +
bg_tasks in one call."""
sid = "s-purge"
await self.bus.session_publish_event(sid, {"e": 1})
await self.bus.inbox_push(sid, {"m": 1})
await self.bus.bg_task_register(sid, "t1", "{}")
await self.bus.session_purge(sid)
self.assertEqual(await self.bus.session_read_events(sid), [])
self.assertEqual(
await self.bus.inbox_drain(sid, max_count=10),
[],
)
self.assertEqual(await self.bus.bg_task_list(sid), {})
async def test_task_cancel_pub_sub(self) -> None:
"""``task_publish_cancel`` → ``task_subscribe_cancel``
round-trip."""
ready = asyncio.Event()
received: list[str] = []
async def _consumer() -> None:
async for tid in self.bus.task_subscribe_cancel(
on_ready=ready.set,
):
received.append(tid)
break
task = asyncio.create_task(_consumer())
await asyncio.wait_for(ready.wait(), timeout=2.0)
await self.bus.task_publish_cancel("task-X")
await asyncio.wait_for(task, timeout=2.0)
self.assertEqual(received, ["task-X"])