Felix Weinberger
330 lines
13 KiB
Python
330 lines
13 KiB
Python
"""Tests for TaskMessageQueue and InMemoryTaskMessageQueue."""
|
|
|
|
from datetime import datetime, timezone
|
|
|
|
import anyio
|
|
import pytest
|
|
|
|
from mcp.shared.experimental.tasks.message_queue import InMemoryTaskMessageQueue, QueuedMessage
|
|
from mcp.shared.experimental.tasks.resolver import Resolver
|
|
from mcp.types import JSONRPCNotification, JSONRPCRequest
|
|
|
|
|
|
@pytest.fixture
|
|
def queue() -> InMemoryTaskMessageQueue:
|
|
return InMemoryTaskMessageQueue()
|
|
|
|
|
|
def make_request(id: int = 1, method: str = "test/method") -> JSONRPCRequest:
|
|
return JSONRPCRequest(jsonrpc="2.0", id=id, method=method)
|
|
|
|
|
|
def make_notification(method: str = "test/notify") -> JSONRPCNotification:
|
|
return JSONRPCNotification(jsonrpc="2.0", method=method)
|
|
|
|
|
|
class TestInMemoryTaskMessageQueue:
|
|
@pytest.mark.anyio
|
|
async def test_enqueue_and_dequeue(self, queue: InMemoryTaskMessageQueue) -> None:
|
|
"""Test basic enqueue and dequeue operations."""
|
|
task_id = "task-1"
|
|
msg = QueuedMessage(type="request", message=make_request())
|
|
|
|
await queue.enqueue(task_id, msg)
|
|
result = await queue.dequeue(task_id)
|
|
|
|
assert result is not None
|
|
assert result.type == "request"
|
|
assert result.message.method == "test/method"
|
|
|
|
@pytest.mark.anyio
|
|
async def test_dequeue_empty_returns_none(self, queue: InMemoryTaskMessageQueue) -> None:
|
|
"""Dequeue from empty queue returns None."""
|
|
result = await queue.dequeue("nonexistent-task")
|
|
assert result is None
|
|
|
|
@pytest.mark.anyio
|
|
async def test_fifo_ordering(self, queue: InMemoryTaskMessageQueue) -> None:
|
|
"""Messages are dequeued in FIFO order."""
|
|
task_id = "task-1"
|
|
|
|
await queue.enqueue(task_id, QueuedMessage(type="request", message=make_request(1, "first")))
|
|
await queue.enqueue(task_id, QueuedMessage(type="request", message=make_request(2, "second")))
|
|
await queue.enqueue(task_id, QueuedMessage(type="request", message=make_request(3, "third")))
|
|
|
|
msg1 = await queue.dequeue(task_id)
|
|
msg2 = await queue.dequeue(task_id)
|
|
msg3 = await queue.dequeue(task_id)
|
|
|
|
assert msg1 is not None and msg1.message.method == "first"
|
|
assert msg2 is not None and msg2.message.method == "second"
|
|
assert msg3 is not None and msg3.message.method == "third"
|
|
|
|
@pytest.mark.anyio
|
|
async def test_separate_queues_per_task(self, queue: InMemoryTaskMessageQueue) -> None:
|
|
"""Each task has its own queue."""
|
|
await queue.enqueue("task-1", QueuedMessage(type="request", message=make_request(1, "task1-msg")))
|
|
await queue.enqueue("task-2", QueuedMessage(type="request", message=make_request(2, "task2-msg")))
|
|
|
|
msg1 = await queue.dequeue("task-1")
|
|
msg2 = await queue.dequeue("task-2")
|
|
|
|
assert msg1 is not None and msg1.message.method == "task1-msg"
|
|
assert msg2 is not None and msg2.message.method == "task2-msg"
|
|
|
|
@pytest.mark.anyio
|
|
async def test_peek_does_not_remove(self, queue: InMemoryTaskMessageQueue) -> None:
|
|
"""Peek returns message without removing it."""
|
|
task_id = "task-1"
|
|
await queue.enqueue(task_id, QueuedMessage(type="request", message=make_request()))
|
|
|
|
peeked = await queue.peek(task_id)
|
|
dequeued = await queue.dequeue(task_id)
|
|
|
|
assert peeked is not None
|
|
assert dequeued is not None
|
|
assert isinstance(peeked.message, JSONRPCRequest)
|
|
assert isinstance(dequeued.message, JSONRPCRequest)
|
|
assert peeked.message.id == dequeued.message.id
|
|
|
|
@pytest.mark.anyio
|
|
async def test_is_empty(self, queue: InMemoryTaskMessageQueue) -> None:
|
|
"""Test is_empty method."""
|
|
task_id = "task-1"
|
|
|
|
assert await queue.is_empty(task_id) is True
|
|
|
|
await queue.enqueue(task_id, QueuedMessage(type="notification", message=make_notification()))
|
|
assert await queue.is_empty(task_id) is False
|
|
|
|
await queue.dequeue(task_id)
|
|
assert await queue.is_empty(task_id) is True
|
|
|
|
@pytest.mark.anyio
|
|
async def test_clear_returns_all_messages(self, queue: InMemoryTaskMessageQueue) -> None:
|
|
"""Clear removes and returns all messages."""
|
|
task_id = "task-1"
|
|
|
|
await queue.enqueue(task_id, QueuedMessage(type="request", message=make_request(1)))
|
|
await queue.enqueue(task_id, QueuedMessage(type="request", message=make_request(2)))
|
|
await queue.enqueue(task_id, QueuedMessage(type="request", message=make_request(3)))
|
|
|
|
messages = await queue.clear(task_id)
|
|
|
|
assert len(messages) == 3
|
|
assert await queue.is_empty(task_id) is True
|
|
|
|
@pytest.mark.anyio
|
|
async def test_clear_empty_queue(self, queue: InMemoryTaskMessageQueue) -> None:
|
|
"""Clear on empty queue returns empty list."""
|
|
messages = await queue.clear("nonexistent")
|
|
assert messages == []
|
|
|
|
@pytest.mark.anyio
|
|
async def test_notification_messages(self, queue: InMemoryTaskMessageQueue) -> None:
|
|
"""Test queuing notification messages."""
|
|
task_id = "task-1"
|
|
msg = QueuedMessage(type="notification", message=make_notification("log/message"))
|
|
|
|
await queue.enqueue(task_id, msg)
|
|
result = await queue.dequeue(task_id)
|
|
|
|
assert result is not None
|
|
assert result.type == "notification"
|
|
assert result.message.method == "log/message"
|
|
|
|
@pytest.mark.anyio
|
|
async def test_message_timestamp(self, queue: InMemoryTaskMessageQueue) -> None:
|
|
"""Messages have timestamps."""
|
|
before = datetime.now(timezone.utc)
|
|
msg = QueuedMessage(type="request", message=make_request())
|
|
after = datetime.now(timezone.utc)
|
|
|
|
assert before <= msg.timestamp <= after
|
|
|
|
@pytest.mark.anyio
|
|
async def test_message_with_resolver(self, queue: InMemoryTaskMessageQueue) -> None:
|
|
"""Messages can have resolvers."""
|
|
task_id = "task-1"
|
|
resolver: Resolver[dict[str, str]] = Resolver()
|
|
|
|
msg = QueuedMessage(
|
|
type="request",
|
|
message=make_request(),
|
|
resolver=resolver,
|
|
original_request_id=42,
|
|
)
|
|
|
|
await queue.enqueue(task_id, msg)
|
|
result = await queue.dequeue(task_id)
|
|
|
|
assert result is not None
|
|
assert result.resolver is resolver
|
|
assert result.original_request_id == 42
|
|
|
|
@pytest.mark.anyio
|
|
async def test_cleanup_specific_task(self, queue: InMemoryTaskMessageQueue) -> None:
|
|
"""Cleanup removes specific task's data."""
|
|
await queue.enqueue("task-1", QueuedMessage(type="request", message=make_request(1)))
|
|
await queue.enqueue("task-2", QueuedMessage(type="request", message=make_request(2)))
|
|
|
|
queue.cleanup("task-1")
|
|
|
|
assert await queue.is_empty("task-1") is True
|
|
assert await queue.is_empty("task-2") is False
|
|
|
|
@pytest.mark.anyio
|
|
async def test_cleanup_all(self, queue: InMemoryTaskMessageQueue) -> None:
|
|
"""Cleanup without task_id removes all data."""
|
|
await queue.enqueue("task-1", QueuedMessage(type="request", message=make_request(1)))
|
|
await queue.enqueue("task-2", QueuedMessage(type="request", message=make_request(2)))
|
|
|
|
queue.cleanup()
|
|
|
|
assert await queue.is_empty("task-1") is True
|
|
assert await queue.is_empty("task-2") is True
|
|
|
|
@pytest.mark.anyio
|
|
async def test_wait_for_message_returns_immediately_if_message_exists(
|
|
self, queue: InMemoryTaskMessageQueue
|
|
) -> None:
|
|
"""wait_for_message returns immediately if queue not empty."""
|
|
task_id = "task-1"
|
|
await queue.enqueue(task_id, QueuedMessage(type="request", message=make_request()))
|
|
|
|
# Should return immediately, not block
|
|
with anyio.fail_after(1):
|
|
await queue.wait_for_message(task_id)
|
|
|
|
@pytest.mark.anyio
|
|
async def test_wait_for_message_blocks_until_message(self, queue: InMemoryTaskMessageQueue) -> None:
|
|
"""wait_for_message blocks until a message is enqueued."""
|
|
task_id = "task-1"
|
|
received = False
|
|
waiter_started = anyio.Event()
|
|
|
|
async def enqueue_when_ready() -> None:
|
|
# Wait until the waiter has started before enqueueing
|
|
await waiter_started.wait()
|
|
await queue.enqueue(task_id, QueuedMessage(type="request", message=make_request()))
|
|
|
|
async def wait_for_msg() -> None:
|
|
nonlocal received
|
|
# Signal that we're about to start waiting
|
|
waiter_started.set()
|
|
await queue.wait_for_message(task_id)
|
|
received = True
|
|
|
|
async with anyio.create_task_group() as tg:
|
|
tg.start_soon(wait_for_msg)
|
|
tg.start_soon(enqueue_when_ready)
|
|
|
|
assert received is True
|
|
|
|
@pytest.mark.anyio
|
|
async def test_notify_message_available_wakes_waiter(self, queue: InMemoryTaskMessageQueue) -> None:
|
|
"""notify_message_available wakes up waiting coroutines."""
|
|
task_id = "task-1"
|
|
notified = False
|
|
waiter_started = anyio.Event()
|
|
|
|
async def notify_when_ready() -> None:
|
|
# Wait until the waiter has started before notifying
|
|
await waiter_started.wait()
|
|
await queue.notify_message_available(task_id)
|
|
|
|
async def wait_for_notification() -> None:
|
|
nonlocal notified
|
|
# Signal that we're about to start waiting
|
|
waiter_started.set()
|
|
await queue.wait_for_message(task_id)
|
|
notified = True
|
|
|
|
async with anyio.create_task_group() as tg:
|
|
tg.start_soon(wait_for_notification)
|
|
tg.start_soon(notify_when_ready)
|
|
|
|
assert notified is True
|
|
|
|
@pytest.mark.anyio
|
|
async def test_peek_empty_queue_returns_none(self, queue: InMemoryTaskMessageQueue) -> None:
|
|
"""Peek on empty queue returns None."""
|
|
result = await queue.peek("nonexistent-task")
|
|
assert result is None
|
|
|
|
@pytest.mark.anyio
|
|
async def test_wait_for_message_double_check_race_condition(self, queue: InMemoryTaskMessageQueue) -> None:
|
|
"""wait_for_message returns early if message arrives after event creation but before wait."""
|
|
task_id = "task-1"
|
|
|
|
# To test the double-check path (lines 223-225), we need a message to arrive
|
|
# after the event is created (line 220) but before event.wait() (line 228).
|
|
# We simulate this by injecting a message before is_empty is called the second time.
|
|
|
|
original_is_empty = queue.is_empty
|
|
call_count = 0
|
|
|
|
async def is_empty_with_injection(tid: str) -> bool:
|
|
nonlocal call_count
|
|
call_count += 1
|
|
if call_count == 2 and tid == task_id:
|
|
# Before second check, inject a message - this simulates a message
|
|
# arriving between event creation and the double-check
|
|
queue._queues[task_id] = [QueuedMessage(type="request", message=make_request())]
|
|
return await original_is_empty(tid)
|
|
|
|
queue.is_empty = is_empty_with_injection # type: ignore[method-assign]
|
|
|
|
# Should return immediately due to double-check finding the message
|
|
with anyio.fail_after(1):
|
|
await queue.wait_for_message(task_id)
|
|
|
|
|
|
class TestResolver:
|
|
@pytest.mark.anyio
|
|
async def test_set_result_and_wait(self) -> None:
|
|
"""Test basic set_result and wait flow."""
|
|
resolver: Resolver[str] = Resolver()
|
|
|
|
resolver.set_result("hello")
|
|
result = await resolver.wait()
|
|
|
|
assert result == "hello"
|
|
assert resolver.done()
|
|
|
|
@pytest.mark.anyio
|
|
async def test_set_exception_and_wait(self) -> None:
|
|
"""Test set_exception raises on wait."""
|
|
resolver: Resolver[str] = Resolver()
|
|
|
|
resolver.set_exception(ValueError("test error"))
|
|
|
|
with pytest.raises(ValueError, match="test error"):
|
|
await resolver.wait()
|
|
|
|
assert resolver.done()
|
|
|
|
@pytest.mark.anyio
|
|
async def test_set_result_when_already_completed_raises(self) -> None:
|
|
"""Test that set_result raises if resolver already completed."""
|
|
resolver: Resolver[str] = Resolver()
|
|
resolver.set_result("first")
|
|
|
|
with pytest.raises(RuntimeError, match="already completed"):
|
|
resolver.set_result("second")
|
|
|
|
@pytest.mark.anyio
|
|
async def test_set_exception_when_already_completed_raises(self) -> None:
|
|
"""Test that set_exception raises if resolver already completed."""
|
|
resolver: Resolver[str] = Resolver()
|
|
resolver.set_result("done")
|
|
|
|
with pytest.raises(RuntimeError, match="already completed"):
|
|
resolver.set_exception(ValueError("too late"))
|
|
|
|
@pytest.mark.anyio
|
|
async def test_done_returns_false_before_completion(self) -> None:
|
|
"""Test done() returns False before any result is set."""
|
|
resolver: Resolver[str] = Resolver()
|
|
assert resolver.done() is False
|