Skip to content

Commit 621081e

Browse files
feat: add ergonomic Client class for testing MCP servers
Add a high-level Client class that wraps ClientSession with transport management, making it easier to test MCP servers. Key changes: - Add Client class with async context manager support - Add InMemoryTransport for in-memory server connections - Migrate test files to use new Client(server) API - Update testing documentation with Client usage examples Usage: async with Client(server) as client: result = await client.call_tool('my_tool', {'arg': 'value'})
1 parent 8adb5bd commit 621081e

33 files changed

+1298
-526
lines changed

docs/testing.md

Lines changed: 12 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,11 @@
11
# Testing MCP Servers
22

3-
If you call yourself a developer, you will want to test your MCP server.
4-
The Python SDK offers the `create_connected_server_and_client_session` function to create a session
5-
using an in-memory transport. I know, I know, the name is too long... We are working on improving it.
3+
The Python SDK provides a `Client` class for testing MCP servers with an in-memory transport.
4+
This makes it easy to write tests without network overhead.
65

7-
Anyway, let's assume you have a simple server with a single tool:
6+
## Basic Usage
7+
8+
Let's assume you have a simple server with a single tool:
89

910
```python title="server.py"
1011
from mcp.server import FastMCP
@@ -40,12 +41,9 @@ To run the below test, you'll need to install the following dependencies:
4041
server - you don't need to use it, but we are spreading the word for best practices.
4142

4243
```python title="test_server.py"
43-
from collections.abc import AsyncGenerator
44-
4544
import pytest
4645
from inline_snapshot import snapshot
47-
from mcp.client.session import ClientSession
48-
from mcp.shared.memory import create_connected_server_and_client_session
46+
from mcp import Client
4947
from mcp.types import CallToolResult, TextContent
5048

5149
from server import app
@@ -57,14 +55,14 @@ def anyio_backend(): # (1)!
5755

5856

5957
@pytest.fixture
60-
async def client_session() -> AsyncGenerator[ClientSession]:
61-
async with create_connected_server_and_client_session(app, raise_exceptions=True) as _session:
62-
yield _session
58+
async def client(): # (2)!
59+
async with Client(app, raise_exceptions=True) as c:
60+
yield c
6361

6462

6563
@pytest.mark.anyio
66-
async def test_call_add_tool(client_session: ClientSession):
67-
result = await client_session.call_tool("add", {"a": 1, "b": 2})
64+
async def test_call_add_tool(client: Client):
65+
result = await client.call_tool("add", {"a": 1, "b": 2})
6866
assert result == snapshot(
6967
CallToolResult(
7068
content=[TextContent(type="text", text="3")],
@@ -74,5 +72,6 @@ async def test_call_add_tool(client_session: ClientSession):
7472
```
7573

7674
1. If you are using `trio`, you should set `"trio"` as the `anyio_backend`. Check more information in the [anyio documentation](https://anyio.readthedocs.io/en/stable/testing.html#specifying-the-backends-to-run-on).
75+
2. The `client` fixture creates a connected client that can be reused across multiple tests.
7776

7877
There you go! You can now extend your tests to cover more scenarios.

examples/fastmcp/weather_structured.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,8 @@
1414

1515
from pydantic import BaseModel, Field
1616

17+
from mcp.client import Client
1718
from mcp.server.fastmcp import FastMCP
18-
from mcp.shared.memory import create_connected_server_and_client_session as client_session
1919

2020
# Create server
2121
mcp = FastMCP("Weather Service")
@@ -157,7 +157,7 @@ async def test() -> None:
157157
print("Testing Weather Service Tools (via MCP protocol)\n")
158158
print("=" * 80)
159159

160-
async with client_session(mcp._mcp_server) as client:
160+
async with Client(mcp) as client:
161161
# Test get_weather
162162
result = await client.call_tool("get_weather", {"city": "London"})
163163
print("\nWeather in London:")

src/mcp/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
from .client.client import Client
12
from .client.session import ClientSession
23
from .client.session_group import ClientSessionGroup
34
from .client.stdio import StdioServerParameters, stdio_client
@@ -66,6 +67,7 @@
6667

6768
__all__ = [
6869
"CallToolRequest",
70+
"Client",
6971
"ClientCapabilities",
7072
"ClientNotification",
7173
"ClientRequest",

src/mcp/client/__init__.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
"""MCP Client module."""
2+
3+
from mcp.client.client import Client
4+
from mcp.client.session import ClientSession
5+
6+
__all__ = [
7+
"Client",
8+
"ClientSession",
9+
]

src/mcp/client/_memory.py

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
"""In-memory transport for testing MCP servers without network overhead."""
2+
3+
from __future__ import annotations
4+
5+
from collections.abc import AsyncGenerator
6+
from contextlib import asynccontextmanager
7+
from typing import Any
8+
9+
import anyio
10+
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
11+
12+
from mcp.server import Server
13+
from mcp.server.fastmcp import FastMCP
14+
from mcp.shared.memory import create_client_server_memory_streams
15+
from mcp.shared.message import SessionMessage
16+
17+
18+
class InMemoryTransport:
19+
"""
20+
In-memory transport for testing MCP servers without network overhead.
21+
22+
This transport starts the server in a background task and provides
23+
streams for client-side communication. The server is automatically
24+
stopped when the context manager exits.
25+
26+
Example:
27+
server = FastMCP("test")
28+
transport = InMemoryTransport(server)
29+
30+
async with transport.connect() as (read_stream, write_stream):
31+
async with ClientSession(read_stream, write_stream) as session:
32+
await session.initialize()
33+
# Use the session...
34+
35+
Or more commonly, use with Client:
36+
async with Client(server) as client:
37+
result = await client.call_tool("my_tool", {...})
38+
"""
39+
40+
def __init__(
41+
self,
42+
server: Server[Any] | FastMCP,
43+
*,
44+
raise_exceptions: bool = False,
45+
) -> None:
46+
"""
47+
Initialize the in-memory transport.
48+
49+
Args:
50+
server: The MCP server to connect to (Server or FastMCP instance)
51+
raise_exceptions: Whether to raise exceptions from the server
52+
"""
53+
self._server = server
54+
self._raise_exceptions = raise_exceptions
55+
56+
@asynccontextmanager
57+
async def connect(
58+
self,
59+
) -> AsyncGenerator[
60+
tuple[
61+
MemoryObjectReceiveStream[SessionMessage | Exception],
62+
MemoryObjectSendStream[SessionMessage],
63+
],
64+
None,
65+
]:
66+
"""
67+
Connect to the server and return streams for communication.
68+
69+
Yields:
70+
A tuple of (read_stream, write_stream) for bidirectional communication
71+
"""
72+
# Unwrap FastMCP to get underlying Server
73+
actual_server: Server[Any]
74+
if isinstance(self._server, FastMCP):
75+
actual_server = self._server._mcp_server # type: ignore[reportPrivateUsage]
76+
else:
77+
actual_server = self._server
78+
79+
async with create_client_server_memory_streams() as (client_streams, server_streams):
80+
client_read, client_write = client_streams
81+
server_read, server_write = server_streams
82+
83+
async with anyio.create_task_group() as tg:
84+
# Start server in background
85+
tg.start_soon(
86+
lambda: actual_server.run(
87+
server_read,
88+
server_write,
89+
actual_server.create_initialization_options(),
90+
raise_exceptions=self._raise_exceptions,
91+
)
92+
)
93+
94+
try:
95+
yield client_read, client_write
96+
finally:
97+
tg.cancel_scope.cancel()

0 commit comments

Comments
 (0)