modelcontextprotocol/python-sdk
Watch
1
0
Fork
You've already forked python-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add ergonomic Client class for testing MCP servers (#1870)

Browse source
This commit is contained in:
Felix Weinberger 2026-01-16 15:49:26 +00:00 • committed by GitHub
commit df039bf97c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 1298 additions and 526 deletions
Download patch file Download diff file Expand all files Collapse all files

21
tests/shared/test_memory.py
View file

@ -1,9 +1,7 @@
import pytest
from typing_extensions import AsyncGenerator
from mcp.client.session import ClientSession
from mcp import Client
from mcp.server import Server
from mcp.shared.memory import create_connected_server_and_client_session
from mcp.types import EmptyResult, Resource
@ -24,18 +22,9 @@ def mcp_server() -> Server:
return server
@pytest.fixture
async def client_connected_to_server(
mcp_server: Server,
) -> AsyncGenerator[ClientSession, None]:
async with create_connected_server_and_client_session(mcp_server) as client_session:
yield client_session
@pytest.mark.anyio
async def test_memory_server_and_client_connection(
client_connected_to_server: ClientSession,
):
async def test_memory_server_and_client_connection(mcp_server: Server):
"""Shows how a client and server can communicate over memory streams."""
response = await client_connected_to_server.send_ping()
assert isinstance(response, EmptyResult)
async with Client(mcp_server) as client:
response = await client.send_ping()
assert isinstance(response, EmptyResult)

50
tests/shared/test_progress_notifications.py
View file

@ -5,15 +5,16 @@ import anyio
import pytest
import mcp.types as types
from mcp.client._memory import InMemoryTransport
from mcp.client.session import ClientSession
from mcp.server import Server
from mcp.server.lowlevel import NotificationOptions
from mcp.server.models import InitializationOptions
from mcp.server.session import ServerSession
from mcp.shared.context import RequestContext
from mcp.shared.memory import create_connected_server_and_client_session
from mcp.shared.message import SessionMessage
from mcp.shared.progress import progress
from mcp.shared.session import BaseSession, RequestResponder, SessionMessage
from mcp.shared.session import BaseSession, RequestResponder
@pytest.mark.anyio
@ -368,25 +369,30 @@ async def test_progress_callback_exception_logging():
# Test with mocked logging
with patch("mcp.shared.session.logging.error", side_effect=mock_log_error):
async with create_connected_server_and_client_session(server) as client_session:
# Send a request with a failing progress callback
result = await client_session.send_request(
types.ClientRequest(
types.CallToolRequest(
method="tools/call",
params=types.CallToolRequestParams(name="progress_tool", arguments={}),
)
),
types.CallToolResult,
progress_callback=failing_progress_callback,
)
transport = InMemoryTransport(server)
async with transport.connect() as (read_stream, write_stream):
async with ClientSession( # pragma: no branch
read_stream=read_stream, write_stream=write_stream
) as session:
await session.initialize()
# Send a request with a failing progress callback
result = await session.send_request(
types.ClientRequest(
types.CallToolRequest(
method="tools/call",
params=types.CallToolRequestParams(name="progress_tool", arguments={}),
)
),
types.CallToolResult,
progress_callback=failing_progress_callback,
)
# Verify the request completed successfully despite the callback failure
assert len(result.content) == 1
content = result.content[0]
assert isinstance(content, types.TextContent)
assert content.text == "progress_result"
# Verify the request completed successfully despite the callback failure
assert len(result.content) == 1
content = result.content[0]
assert isinstance(content, types.TextContent)
assert content.text == "progress_result"
# Check that a warning was logged for the progress callback exception
assert len(logged_errors) > 0
assert any("Progress callback raised an exception" in warning for warning in logged_errors)
# Check that a warning was logged for the progress callback exception
assert len(logged_errors) > 0
assert any("Progress callback raised an exception" in warning for warning in logged_errors)

72
tests/shared/test_session.py
View file

@ -1,14 +1,14 @@
from collections.abc import AsyncGenerator
from typing import Any
import anyio
import pytest
import mcp.types as types
from mcp.client._memory import InMemoryTransport
from mcp.client.session import ClientSession
from mcp.server.lowlevel.server import Server
from mcp.shared.exceptions import McpError
from mcp.shared.memory import create_client_server_memory_streams, create_connected_server_and_client_session
from mcp.shared.memory import create_client_server_memory_streams
from mcp.shared.message import SessionMessage
from mcp.types import (
CancelledNotification,
@ -30,25 +30,20 @@ def mcp_server() -> Server:
return Server(name="test server")
@pytest.fixture
async def client_connected_to_server(
mcp_server: Server,
) -> AsyncGenerator[ClientSession, None]:
async with create_connected_server_and_client_session(mcp_server) as client_session:
yield client_session
@pytest.mark.anyio
async def test_in_flight_requests_cleared_after_completion(
client_connected_to_server: ClientSession,
):
async def test_in_flight_requests_cleared_after_completion(mcp_server: Server):
"""Verify that _in_flight is empty after all requests complete."""
# Send a request and wait for response
response = await client_connected_to_server.send_ping()
assert isinstance(response, EmptyResult)
transport = InMemoryTransport(mcp_server)
async with transport.connect() as (read_stream, write_stream):
async with ClientSession(read_stream=read_stream, write_stream=write_stream) as session:
await session.initialize()
# Verify _in_flight is empty
assert len(client_connected_to_server._in_flight) == 0
# Send a request and wait for response
response = await session.send_ping()
assert isinstance(response, EmptyResult)
# Verify _in_flight is empty
assert len(session._in_flight) == 0
@pytest.mark.anyio
@ -88,10 +83,10 @@ async def test_request_cancellation():
return server
async def make_request(client_session: ClientSession):
async def make_request(session: ClientSession):
nonlocal ev_cancelled
try:
await client_session.send_request(
await session.send_request(
ClientRequest(
types.CallToolRequest(
params=types.CallToolRequestParams(name="slow_tool", arguments={}),
@ -105,28 +100,31 @@ async def test_request_cancellation():
assert "Request cancelled" in str(e)
ev_cancelled.set()
async with create_connected_server_and_client_session(make_server()) as client_session:
async with anyio.create_task_group() as tg:
tg.start_soon(make_request, client_session)
transport = InMemoryTransport(make_server())
async with transport.connect() as (read_stream, write_stream):
async with ClientSession(read_stream=read_stream, write_stream=write_stream) as session:
await session.initialize()
async with anyio.create_task_group() as tg: # pragma: no branch
tg.start_soon(make_request, session)
# Wait for the request to be in-flight
with anyio.fail_after(1): # Timeout after 1 second
await ev_tool_called.wait()
# Wait for the request to be in-flight
with anyio.fail_after(1): # Timeout after 1 second
await ev_tool_called.wait()
# Send cancellation notification
assert request_id is not None
await client_session.send_notification(
ClientNotification(
CancelledNotification(
params=CancelledNotificationParams(request_id=request_id),
# Send cancellation notification
assert request_id is not None
await session.send_notification(
ClientNotification(
CancelledNotification(
params=CancelledNotificationParams(request_id=request_id),
)
)
)
)
# Give cancellation time to process
# TODO(Marcelo): Drop the pragma once https://github.com/coveragepy/coveragepy/issues/1987 is fixed.
with anyio.fail_after(1): # pragma: no cover
await ev_cancelled.wait()
# Give cancellation time to process
# TODO(Marcelo): Drop the pragma once https://github.com/coveragepy/coveragepy/issues/1987 is fixed.
with anyio.fail_after(1): # pragma: no cover
await ev_cancelled.wait()
@pytest.mark.anyio