-
Notifications
You must be signed in to change notification settings - Fork 3.3k
Expand file tree
/
Copy pathtest_188_concurrency.py
More file actions
85 lines (72 loc) · 2.73 KB
/
test_188_concurrency.py
File metadata and controls
85 lines (72 loc) · 2.73 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
import anyio
import pytest
from mcp import Client
from mcp.server.mcpserver import MCPServer
@pytest.mark.anyio
async def test_messages_are_executed_concurrently_tools():
server = MCPServer("test")
event = anyio.Event()
tool_started = anyio.Event()
call_order: list[str] = []
@server.tool("sleep")
async def sleep_tool():
call_order.append("waiting_for_event")
tool_started.set()
await event.wait()
call_order.append("tool_end")
return "done"
@server.tool("trigger")
async def trigger():
# Wait for tool to start before setting the event
await tool_started.wait()
call_order.append("trigger_started")
event.set()
call_order.append("trigger_end")
return "slow"
async with Client(server) as client_session:
# First tool will wait on event, second will set it
async with anyio.create_task_group() as tg:
# Start the tool first (it will wait on event)
tg.start_soon(client_session.call_tool, "sleep")
# Then the trigger tool will set the event to allow the first tool to continue
await client_session.call_tool("trigger")
# Verify that both ran concurrently
assert call_order == [
"waiting_for_event",
"trigger_started",
"trigger_end",
"tool_end",
], f"Expected concurrent execution, but got: {call_order}"
@pytest.mark.anyio
async def test_messages_are_executed_concurrently_tools_and_resources():
server = MCPServer("test")
event = anyio.Event()
tool_started = anyio.Event()
call_order: list[str] = []
@server.tool("sleep")
async def sleep_tool():
call_order.append("waiting_for_event")
tool_started.set()
await event.wait()
call_order.append("tool_end")
return "done"
@server.resource("slow://slow_resource")
async def slow_resource():
# Wait for tool to start before setting the event
await tool_started.wait()
event.set()
call_order.append("resource_end")
return "slow"
async with Client(server) as client_session:
# First tool will wait on event, second will set it
async with anyio.create_task_group() as tg:
# Start the tool first (it will wait on event)
tg.start_soon(client_session.call_tool, "sleep")
# Then the resource (it will set the event)
tg.start_soon(client_session.read_resource, "slow://slow_resource")
# Verify that both ran concurrently
assert call_order == [
"waiting_for_event",
"resource_end",
"tool_end",
], f"Expected concurrent execution, but got: {call_order}"