refactor: use Client class in tests (#1900)

This commit is contained in:
Felix Weinberger 2026-01-17 08:46:58 +00:00 • committed by GitHub
commit dcc9b4f8e0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 163 additions and 227 deletions

View file

@ -3,8 +3,7 @@ from collections.abc import Callable
import pytest import pytest
import mcp.types as types import mcp.types as types
from mcp.client._memory import InMemoryTransport from mcp import Client
from mcp.client.session import ClientSession
from mcp.server import Server from mcp.server import Server
from mcp.server.fastmcp import FastMCP from mcp.server.fastmcp import FastMCP
from mcp.types import ListToolsRequest, ListToolsResult from mcp.types import ListToolsRequest, ListToolsResult
@ -66,49 +65,43 @@ async def test_list_methods_params_parameter(
See: https://modelcontextprotocol.io/specification/2025-03-26/server/utilities/pagination#request-format See: https://modelcontextprotocol.io/specification/2025-03-26/server/utilities/pagination#request-format
""" """
transport = InMemoryTransport(full_featured_server) async with Client(full_featured_server) as client:
async with transport.connect() as (read_stream, write_stream): spies = stream_spy()
async with ClientSession(read_stream, write_stream) as session:
await session.initialize()
spies = stream_spy()
# Test without params (omitted) # Test without params (omitted)
method = getattr(session, method_name) method = getattr(client, method_name)
_ = await method() _ = await method()
requests = spies.get_client_requests(method=request_method) requests = spies.get_client_requests(method=request_method)
assert len(requests) == 1 assert len(requests) == 1
assert requests[0].params is None assert requests[0].params is None
spies.clear() spies.clear()
# Test with params containing cursor # Test with params containing cursor
_ = await method(params=types.PaginatedRequestParams(cursor="from_params")) _ = await method(params=types.PaginatedRequestParams(cursor="from_params"))
requests = spies.get_client_requests(method=request_method) requests = spies.get_client_requests(method=request_method)
assert len(requests) == 1 assert len(requests) == 1
assert requests[0].params is not None assert requests[0].params is not None
assert requests[0].params["cursor"] == "from_params" assert requests[0].params["cursor"] == "from_params"
spies.clear() spies.clear()
# Test with empty params # Test with empty params
_ = await method(params=types.PaginatedRequestParams()) _ = await method(params=types.PaginatedRequestParams())
requests = spies.get_client_requests(method=request_method) requests = spies.get_client_requests(method=request_method)
assert len(requests) == 1 assert len(requests) == 1
# Empty params means no cursor # Empty params means no cursor
assert requests[0].params is None or "cursor" not in requests[0].params assert requests[0].params is None or "cursor" not in requests[0].params
async def test_list_tools_with_strict_server_validation( async def test_list_tools_with_strict_server_validation(
full_featured_server: FastMCP, full_featured_server: FastMCP,
): ):
"""Test pagination with a server that validates request format strictly.""" """Test pagination with a server that validates request format strictly."""
transport = InMemoryTransport(full_featured_server) async with Client(full_featured_server) as client:
async with transport.connect() as (read_stream, write_stream): result = await client.list_tools(params=types.PaginatedRequestParams())
async with ClientSession(read_stream, write_stream) as session: assert isinstance(result, ListToolsResult)
await session.initialize() assert len(result.tools) > 0
result = await session.list_tools(params=types.PaginatedRequestParams())
assert isinstance(result, ListToolsResult)
assert len(result.tools) > 0
async def test_list_tools_with_lowlevel_server(): async def test_list_tools_with_lowlevel_server():
@ -129,13 +122,9 @@ async def test_list_tools_with_lowlevel_server():
] ]
) )
transport = InMemoryTransport(server) async with Client(server) as client:
async with transport.connect() as (read_stream, write_stream): result = await client.list_tools(params=types.PaginatedRequestParams())
async with ClientSession(read_stream, write_stream) as session: assert result.tools[0].description == "cursor=None"
await session.initialize()
result = await session.list_tools(params=types.PaginatedRequestParams()) result = await client.list_tools(params=types.PaginatedRequestParams(cursor="page2"))
assert result.tools[0].description == "cursor=None" assert result.tools[0].description == "cursor=page2"
result = await session.list_tools(params=types.PaginatedRequestParams(cursor="page2"))
assert result.tools[0].description == "cursor=page2"

View file

@ -2,6 +2,7 @@
import pytest import pytest
from mcp import Client
from mcp.client._memory import InMemoryTransport from mcp.client._memory import InMemoryTransport
from mcp.server import Server from mcp.server import Server
from mcp.server.fastmcp import FastMCP from mcp.server.fastmcp import FastMCP
@ -69,42 +70,26 @@ async def test_with_fastmcp(fastmcp_server: FastMCP):
async def test_server_is_running(fastmcp_server: FastMCP): async def test_server_is_running(fastmcp_server: FastMCP):
"""Test that the server is running and responding to requests.""" """Test that the server is running and responding to requests."""
from mcp.client.session import ClientSession async with Client(fastmcp_server) as client:
assert client.server_capabilities is not None
transport = InMemoryTransport(fastmcp_server)
async with transport.connect() as (read_stream, write_stream):
async with ClientSession(read_stream, write_stream) as session:
result = await session.initialize()
assert result is not None
assert result.server_info.name == "test"
async def test_list_tools(fastmcp_server: FastMCP): async def test_list_tools(fastmcp_server: FastMCP):
"""Test listing tools through the transport.""" """Test listing tools through the transport."""
from mcp.client.session import ClientSession async with Client(fastmcp_server) as client:
tools_result = await client.list_tools()
transport = InMemoryTransport(fastmcp_server) assert len(tools_result.tools) > 0
async with transport.connect() as (read_stream, write_stream): tool_names = [t.name for t in tools_result.tools]
async with ClientSession(read_stream, write_stream) as session: assert "greet" in tool_names
await session.initialize()
tools_result = await session.list_tools()
assert len(tools_result.tools) > 0
tool_names = [t.name for t in tools_result.tools]
assert "greet" in tool_names
async def test_call_tool(fastmcp_server: FastMCP): async def test_call_tool(fastmcp_server: FastMCP):
"""Test calling a tool through the transport.""" """Test calling a tool through the transport."""
from mcp.client.session import ClientSession async with Client(fastmcp_server) as client:
result = await client.call_tool("greet", {"name": "World"})
transport = InMemoryTransport(fastmcp_server) assert result is not None
async with transport.connect() as (read_stream, write_stream): assert len(result.content) > 0
async with ClientSession(read_stream, write_stream) as session: assert "Hello, World!" in str(result.content[0])
await session.initialize()
result = await session.call_tool("greet", {"name": "World"})
assert result is not None
assert len(result.content) > 0
assert "Hello, World!" in str(result.content[0])
async def test_raise_exceptions(fastmcp_server: FastMCP): async def test_raise_exceptions(fastmcp_server: FastMCP):

View file

@ -6,8 +6,7 @@ import anyio
import pytest import pytest
import mcp.types as types import mcp.types as types
from mcp.client._memory import InMemoryTransport from mcp import Client
from mcp.client.session import ClientSession
from mcp.server.lowlevel.server import Server from mcp.server.lowlevel.server import Server
from mcp.shared.exceptions import McpError from mcp.shared.exceptions import McpError
from mcp.types import ( from mcp.types import (
@ -55,61 +54,50 @@ async def test_server_remains_functional_after_cancel():
return [types.TextContent(type="text", text=f"Call number: {call_count}")] return [types.TextContent(type="text", text=f"Call number: {call_count}")]
raise ValueError(f"Unknown tool: {name}") # pragma: no cover raise ValueError(f"Unknown tool: {name}") # pragma: no cover
transport = InMemoryTransport(server) async with Client(server) as client:
async with transport.connect() as (read_stream, write_stream): # First request (will be cancelled)
async with ClientSession(read_stream, write_stream) as client: async def first_request():
await client.initialize() try:
await client.session.send_request(
# First request (will be cancelled) ClientRequest(
async def first_request(): CallToolRequest(
try: params=CallToolRequestParams(name="test_tool", arguments={}),
await client.send_request(
ClientRequest(
CallToolRequest(
params=CallToolRequestParams(name="test_tool", arguments={}),
)
),
CallToolResult,
)
pytest.fail("First request should have been cancelled") # pragma: no cover
except McpError:
pass # Expected
# Start first request
async with anyio.create_task_group() as tg:
tg.start_soon(first_request)
# Wait for it to start
await ev_first_call.wait()
# Cancel it
assert first_request_id is not None
await client.send_notification(
ClientNotification(
CancelledNotification(
params=CancelledNotificationParams(
request_id=first_request_id,
reason="Testing server recovery",
),
) )
),
CallToolResult,
)
pytest.fail("First request should have been cancelled") # pragma: no cover
except McpError:
pass # Expected
# Start first request
async with anyio.create_task_group() as tg:
tg.start_soon(first_request)
# Wait for it to start
await ev_first_call.wait()
# Cancel it
assert first_request_id is not None
await client.session.send_notification(
ClientNotification(
CancelledNotification(
params=CancelledNotificationParams(
request_id=first_request_id,
reason="Testing server recovery",
),
) )
) )
# Second request (should work normally)
result = await client.send_request(
ClientRequest(
CallToolRequest(
params=CallToolRequestParams(name="test_tool", arguments={}),
)
),
CallToolResult,
) )
# Verify second request completed successfully # Second request (should work normally)
assert len(result.content) == 1 result = await client.call_tool("test_tool", {})
# Type narrowing for pyright
content = result.content[0] # Verify second request completed successfully
assert content.type == "text" assert len(result.content) == 1
assert isinstance(content, types.TextContent) # Type narrowing for pyright
assert content.text == "Call number: 2" content = result.content[0]
assert call_count == 2 assert content.type == "text"
assert isinstance(content, types.TextContent)
assert content.text == "Call number: 2"
assert call_count == 2

View file

@ -5,7 +5,7 @@ import anyio
import pytest import pytest
import mcp.types as types import mcp.types as types
from mcp.client._memory import InMemoryTransport from mcp import Client
from mcp.client.session import ClientSession from mcp.client.session import ClientSession
from mcp.server import Server from mcp.server import Server
from mcp.server.lowlevel import NotificationOptions from mcp.server.lowlevel import NotificationOptions
@ -369,30 +369,20 @@ async def test_progress_callback_exception_logging():
# Test with mocked logging # Test with mocked logging
with patch("mcp.shared.session.logging.error", side_effect=mock_log_error): with patch("mcp.shared.session.logging.error", side_effect=mock_log_error):
transport = InMemoryTransport(server) async with Client(server) as client:
async with transport.connect() as (read_stream, write_stream): # Call tool with a failing progress callback
async with ClientSession( # pragma: no branch result = await client.call_tool(
read_stream=read_stream, write_stream=write_stream "progress_tool",
) as session: arguments={},
await session.initialize() progress_callback=failing_progress_callback,
# 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 # Verify the request completed successfully despite the callback failure
assert len(result.content) == 1 assert len(result.content) == 1
content = result.content[0] content = result.content[0]
assert isinstance(content, types.TextContent) assert isinstance(content, types.TextContent)
assert content.text == "progress_result" assert content.text == "progress_result"
# Check that a warning was logged for the progress callback exception # Check that a warning was logged for the progress callback exception
assert len(logged_errors) > 0 assert len(logged_errors) > 0
assert any("Progress callback raised an exception" in warning for warning in logged_errors) assert any("Progress callback raised an exception" in warning for warning in logged_errors)

View file

@ -4,7 +4,7 @@ import anyio
import pytest import pytest
import mcp.types as types import mcp.types as types
from mcp.client._memory import InMemoryTransport from mcp import Client
from mcp.client.session import ClientSession from mcp.client.session import ClientSession
from mcp.server.lowlevel.server import Server from mcp.server.lowlevel.server import Server
from mcp.shared.exceptions import McpError from mcp.shared.exceptions import McpError
@ -25,68 +25,55 @@ from mcp.types import (
) )
@pytest.fixture
def mcp_server() -> Server:
return Server(name="test server")
@pytest.mark.anyio @pytest.mark.anyio
async def test_in_flight_requests_cleared_after_completion(mcp_server: Server): async def test_in_flight_requests_cleared_after_completion():
"""Verify that _in_flight is empty after all requests complete.""" """Verify that _in_flight is empty after all requests complete."""
transport = InMemoryTransport(mcp_server) server = Server(name="test server")
async with transport.connect() as (read_stream, write_stream): async with Client(server) as client:
async with ClientSession(read_stream=read_stream, write_stream=write_stream) as session: # Send a request and wait for response
await session.initialize() response = await client.send_ping()
assert isinstance(response, EmptyResult)
# Send a request and wait for response # Verify _in_flight is empty
response = await session.send_ping() assert len(client.session._in_flight) == 0
assert isinstance(response, EmptyResult)
# Verify _in_flight is empty
assert len(session._in_flight) == 0
@pytest.mark.anyio @pytest.mark.anyio
async def test_request_cancellation(): async def test_request_cancellation():
"""Test that requests can be cancelled while in-flight.""" """Test that requests can be cancelled while in-flight."""
# The tool is already registered in the fixture
ev_tool_called = anyio.Event() ev_tool_called = anyio.Event()
ev_cancelled = anyio.Event() ev_cancelled = anyio.Event()
request_id = None request_id = None
# Start the request in a separate task so we can cancel it # Create a server with a slow tool
def make_server() -> Server: server = Server(name="TestSessionServer")
server = Server(name="TestSessionServer")
# Register the tool handler # Register the tool handler
@server.call_tool() @server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any] | None) -> list[TextContent]: async def handle_call_tool(name: str, arguments: dict[str, Any] | None) -> list[TextContent]:
nonlocal request_id, ev_tool_called nonlocal request_id, ev_tool_called
if name == "slow_tool": if name == "slow_tool":
request_id = server.request_context.request_id request_id = server.request_context.request_id
ev_tool_called.set() ev_tool_called.set()
await anyio.sleep(10) # Long enough to ensure we can cancel await anyio.sleep(10) # Long enough to ensure we can cancel
return [] # pragma: no cover return [] # pragma: no cover
raise ValueError(f"Unknown tool: {name}") # pragma: no cover raise ValueError(f"Unknown tool: {name}") # pragma: no cover
# Register the tool so it shows up in list_tools # Register the tool so it shows up in list_tools
@server.list_tools() @server.list_tools()
async def handle_list_tools() -> list[types.Tool]: async def handle_list_tools() -> list[types.Tool]:
return [ return [
types.Tool( types.Tool(
name="slow_tool", name="slow_tool",
description="A slow tool that takes 10 seconds to complete", description="A slow tool that takes 10 seconds to complete",
input_schema={}, input_schema={},
) )
] ]
return server async def make_request(client: Client):
async def make_request(session: ClientSession):
nonlocal ev_cancelled nonlocal ev_cancelled
try: try:
await session.send_request( await client.session.send_request(
ClientRequest( ClientRequest(
types.CallToolRequest( types.CallToolRequest(
params=types.CallToolRequestParams(name="slow_tool", arguments={}), params=types.CallToolRequestParams(name="slow_tool", arguments={}),
@ -100,31 +87,28 @@ async def test_request_cancellation():
assert "Request cancelled" in str(e) assert "Request cancelled" in str(e)
ev_cancelled.set() ev_cancelled.set()
transport = InMemoryTransport(make_server()) async with Client(server) as client:
async with transport.connect() as (read_stream, write_stream): async with anyio.create_task_group() as tg: # pragma: no branch
async with ClientSession(read_stream=read_stream, write_stream=write_stream) as session: tg.start_soon(make_request, client)
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 # Wait for the request to be in-flight
with anyio.fail_after(1): # Timeout after 1 second with anyio.fail_after(1): # Timeout after 1 second
await ev_tool_called.wait() await ev_tool_called.wait()
# Send cancellation notification # Send cancellation notification
assert request_id is not None assert request_id is not None
await session.send_notification( await client.session.send_notification(
ClientNotification( ClientNotification(
CancelledNotification( CancelledNotification(
params=CancelledNotificationParams(request_id=request_id), params=CancelledNotificationParams(request_id=request_id),
)
) )
) )
)
# Give cancellation time to process # Give cancellation time to process
# TODO(Marcelo): Drop the pragma once https://github.com/coveragepy/coveragepy/issues/1987 is fixed. # TODO(Marcelo): Drop the pragma once https://github.com/coveragepy/coveragepy/issues/1987 is fixed.
with anyio.fail_after(1): # pragma: no cover with anyio.fail_after(1): # pragma: no cover
await ev_cancelled.wait() await ev_cancelled.wait()
@pytest.mark.anyio @pytest.mark.anyio