test: convert test_integration.py to in-memory transport (fix flaky) (#2277)

This commit is contained in:
Max Isbey 2026-03-18 15:25:11 +00:00 • committed by GitHub
commit 7826ade12b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 231 additions and 515 deletions

View file

@ -474,7 +474,7 @@ class ServerSession(
related_request_id, related_request_id,
) )
async def send_resource_list_changed(self) -> None: # pragma: no cover async def send_resource_list_changed(self) -> None:
"""Send a resource list changed notification.""" """Send a resource list changed notification."""
await self.send_notification(types.ResourceListChangedNotification()) await self.send_notification(types.ResourceListChangedNotification())

View file

@ -1,7 +1,7 @@
"""Integration tests for MCPServer server functionality. """Integration tests for MCPServer server functionality.
These tests validate the proper functioning of MCPServer features using focused, These tests validate the proper functioning of MCPServer features using focused,
single-feature servers across different transports (SSE and StreamableHTTP). single-feature example servers over an in-memory transport.
""" """
# TODO(Marcelo): The `examples` package is not being imported as package. We need to solve this. # TODO(Marcelo): The `examples` package is not being imported as package. We need to solve this.
# pyright: reportUnknownMemberType=false # pyright: reportUnknownMemberType=false
@ -10,12 +10,8 @@ single-feature servers across different transports (SSE and StreamableHTTP).
# pyright: reportUnknownArgumentType=false # pyright: reportUnknownArgumentType=false
import json import json
import multiprocessing
import socket
from collections.abc import Generator
import pytest import pytest
import uvicorn
from inline_snapshot import snapshot from inline_snapshot import snapshot
from examples.snippets.servers import ( from examples.snippets.servers import (
@ -30,9 +26,8 @@ from examples.snippets.servers import (
structured_output, structured_output,
tool_progress, tool_progress,
) )
from mcp.client import Client
from mcp.client.session import ClientSession from mcp.client.session import ClientSession
from mcp.client.sse import sse_client
from mcp.client.streamable_http import streamable_http_client
from mcp.shared._context import RequestContext from mcp.shared._context import RequestContext
from mcp.shared.session import RequestResponder from mcp.shared.session import RequestResponder
from mcp.types import ( from mcp.types import (
@ -42,7 +37,6 @@ from mcp.types import (
ElicitRequestParams, ElicitRequestParams,
ElicitResult, ElicitResult,
GetPromptResult, GetPromptResult,
InitializeResult,
LoggingMessageNotification, LoggingMessageNotification,
LoggingMessageNotificationParams, LoggingMessageNotificationParams,
NotificationParams, NotificationParams,
@ -58,7 +52,8 @@ from mcp.types import (
TextResourceContents, TextResourceContents,
ToolListChangedNotification, ToolListChangedNotification,
) )
from tests.test_helpers import wait_for_server
pytestmark = pytest.mark.anyio
class NotificationCollector: class NotificationCollector:
@ -85,105 +80,6 @@ class NotificationCollector:
self.tool_notifications.append(message.params) self.tool_notifications.append(message.params)
# Common fixtures
@pytest.fixture
def server_port() -> int:
"""Get a free port for testing."""
with socket.socket() as s:
s.bind(("127.0.0.1", 0))
return s.getsockname()[1]
@pytest.fixture
def server_url(server_port: int) -> str:
"""Get the server URL for testing."""
return f"http://127.0.0.1:{server_port}"
def run_server_with_transport(module_name: str, port: int, transport: str) -> None: # pragma: no cover
"""Run server with specified transport."""
# Get the MCP instance based on module name
if module_name == "basic_tool":
mcp = basic_tool.mcp
elif module_name == "basic_resource":
mcp = basic_resource.mcp
elif module_name == "basic_prompt":
mcp = basic_prompt.mcp
elif module_name == "tool_progress":
mcp = tool_progress.mcp
elif module_name == "sampling":
mcp = sampling.mcp
elif module_name == "elicitation":
mcp = elicitation.mcp
elif module_name == "completion":
mcp = completion.mcp
elif module_name == "notifications":
mcp = notifications.mcp
elif module_name == "mcpserver_quickstart":
mcp = mcpserver_quickstart.mcp
elif module_name == "structured_output":
mcp = structured_output.mcp
else:
raise ImportError(f"Unknown module: {module_name}")
# Create app based on transport type
if transport == "sse":
app = mcp.sse_app()
elif transport == "streamable-http":
app = mcp.streamable_http_app()
else:
raise ValueError(f"Invalid transport for test server: {transport}")
server = uvicorn.Server(config=uvicorn.Config(app=app, host="127.0.0.1", port=port, log_level="error"))
print(f"Starting {transport} server on port {port}")
server.run()
@pytest.fixture
def server_transport(request: pytest.FixtureRequest, server_port: int) -> Generator[str, None, None]:
"""Start server in a separate process with specified MCP instance and transport.
Args:
request: pytest request with param tuple of (module_name, transport)
server_port: Port to run the server on
Yields:
str: The transport type ('sse' or 'streamable_http')
"""
module_name, transport = request.param
proc = multiprocessing.Process(
target=run_server_with_transport,
args=(module_name, server_port, transport),
daemon=True,
)
proc.start()
# Wait for server to be ready
wait_for_server(server_port)
yield transport
proc.kill()
proc.join(timeout=2)
if proc.is_alive(): # pragma: no cover
print("Server process failed to terminate")
# Helper function to create client based on transport
def create_client_for_transport(transport: str, server_url: str):
"""Create the appropriate client context manager based on transport type."""
if transport == "sse":
endpoint = f"{server_url}/sse"
return sse_client(endpoint)
elif transport == "streamable-http":
endpoint = f"{server_url}/mcp"
return streamable_http_client(endpoint)
else: # pragma: no cover
raise ValueError(f"Invalid transport: {transport}")
# Callback functions for testing
async def sampling_callback( async def sampling_callback(
context: RequestContext[ClientSession], params: CreateMessageRequestParams context: RequestContext[ClientSession], params: CreateMessageRequestParams
) -> CreateMessageResult: ) -> CreateMessageResult:
@ -210,74 +106,40 @@ async def elicitation_callback(context: RequestContext[ClientSession], params: E
return ElicitResult(action="decline") return ElicitResult(action="decline")
# Test basic tools async def test_basic_tools() -> None:
@pytest.mark.anyio
@pytest.mark.parametrize(
"server_transport",
[
("basic_tool", "sse"),
("basic_tool", "streamable-http"),
],
indirect=True,
)
async def test_basic_tools(server_transport: str, server_url: str) -> None:
"""Test basic tool functionality.""" """Test basic tool functionality."""
transport = server_transport async with Client(basic_tool.mcp) as client:
client_cm = create_client_for_transport(transport, server_url) assert client.server_capabilities is not None
assert client.server_capabilities.tools is not None
async with client_cm as (read_stream, write_stream):
async with ClientSession(read_stream, write_stream) as session:
# Test initialization
result = await session.initialize()
assert isinstance(result, InitializeResult)
assert result.server_info.name == "Tool Example"
assert result.capabilities.tools is not None
# Test sum tool # Test sum tool
tool_result = await session.call_tool("sum", {"a": 5, "b": 3}) tool_result = await client.call_tool("sum", {"a": 5, "b": 3})
assert len(tool_result.content) == 1 assert len(tool_result.content) == 1
assert isinstance(tool_result.content[0], TextContent) assert isinstance(tool_result.content[0], TextContent)
assert tool_result.content[0].text == "8" assert tool_result.content[0].text == "8"
# Test weather tool # Test weather tool
weather_result = await session.call_tool("get_weather", {"city": "London"}) weather_result = await client.call_tool("get_weather", {"city": "London"})
assert len(weather_result.content) == 1 assert len(weather_result.content) == 1
assert isinstance(weather_result.content[0], TextContent) assert isinstance(weather_result.content[0], TextContent)
assert "Weather in London: 22degreesC" in weather_result.content[0].text assert "Weather in London: 22degreesC" in weather_result.content[0].text
# Test resources async def test_basic_resources() -> None:
@pytest.mark.anyio
@pytest.mark.parametrize(
"server_transport",
[
("basic_resource", "sse"),
("basic_resource", "streamable-http"),
],
indirect=True,
)
async def test_basic_resources(server_transport: str, server_url: str) -> None:
"""Test basic resource functionality.""" """Test basic resource functionality."""
transport = server_transport async with Client(basic_resource.mcp) as client:
client_cm = create_client_for_transport(transport, server_url) assert client.server_capabilities is not None
assert client.server_capabilities.resources is not None
async with client_cm as (read_stream, write_stream):
async with ClientSession(read_stream, write_stream) as session:
# Test initialization
result = await session.initialize()
assert isinstance(result, InitializeResult)
assert result.server_info.name == "Resource Example"
assert result.capabilities.resources is not None
# Test document resource # Test document resource
doc_content = await session.read_resource("file://documents/readme") doc_content = await client.read_resource("file://documents/readme")
assert isinstance(doc_content, ReadResourceResult) assert isinstance(doc_content, ReadResourceResult)
assert len(doc_content.contents) == 1 assert len(doc_content.contents) == 1
assert isinstance(doc_content.contents[0], TextResourceContents) assert isinstance(doc_content.contents[0], TextResourceContents)
assert "Content of readme" in doc_content.contents[0].text assert "Content of readme" in doc_content.contents[0].text
# Test settings resource # Test settings resource
settings_content = await session.read_resource("config://settings") settings_content = await client.read_resource("config://settings")
assert isinstance(settings_content, ReadResourceResult) assert isinstance(settings_content, ReadResourceResult)
assert len(settings_content.contents) == 1 assert len(settings_content.contents) == 1
assert isinstance(settings_content.contents[0], TextResourceContents) assert isinstance(settings_content.contents[0], TextResourceContents)
@ -286,35 +148,18 @@ async def test_basic_resources(server_transport: str, server_url: str) -> None:
assert settings_json["language"] == "en" assert settings_json["language"] == "en"
# Test prompts async def test_basic_prompts() -> None:
@pytest.mark.anyio
@pytest.mark.parametrize(
"server_transport",
[
("basic_prompt", "sse"),
("basic_prompt", "streamable-http"),
],
indirect=True,
)
async def test_basic_prompts(server_transport: str, server_url: str) -> None:
"""Test basic prompt functionality.""" """Test basic prompt functionality."""
transport = server_transport async with Client(basic_prompt.mcp) as client:
client_cm = create_client_for_transport(transport, server_url) assert client.server_capabilities is not None
assert client.server_capabilities.prompts is not None
async with client_cm as (read_stream, write_stream):
async with ClientSession(read_stream, write_stream) as session:
# Test initialization
result = await session.initialize()
assert isinstance(result, InitializeResult)
assert result.server_info.name == "Prompt Example"
assert result.capabilities.prompts is not None
# Test review_code prompt # Test review_code prompt
prompts = await session.list_prompts() prompts = await client.list_prompts()
review_prompt = next((p for p in prompts.prompts if p.name == "review_code"), None) review_prompt = next((p for p in prompts.prompts if p.name == "review_code"), None)
assert review_prompt is not None assert review_prompt is not None
prompt_result = await session.get_prompt("review_code", {"code": "def hello():\n print('Hello')"}) prompt_result = await client.get_prompt("review_code", {"code": "def hello():\n print('Hello')"})
assert isinstance(prompt_result, GetPromptResult) assert isinstance(prompt_result, GetPromptResult)
assert len(prompt_result.messages) == 1 assert len(prompt_result.messages) == 1
assert isinstance(prompt_result.messages[0].content, TextContent) assert isinstance(prompt_result.messages[0].content, TextContent)
@ -322,7 +167,7 @@ async def test_basic_prompts(server_transport: str, server_url: str) -> None:
assert "def hello():" in prompt_result.messages[0].content.text assert "def hello():" in prompt_result.messages[0].content.text
# Test debug_error prompt # Test debug_error prompt
debug_result = await session.get_prompt( debug_result = await client.get_prompt(
"debug_error", {"error": "TypeError: 'NoneType' object is not subscriptable"} "debug_error", {"error": "TypeError: 'NoneType' object is not subscriptable"}
) )
assert isinstance(debug_result, GetPromptResult) assert isinstance(debug_result, GetPromptResult)
@ -338,19 +183,8 @@ async def test_basic_prompts(server_transport: str, server_url: str) -> None:
assert "I'll help debug that" in debug_result.messages[2].content.text assert "I'll help debug that" in debug_result.messages[2].content.text
# Test progress reporting async def test_tool_progress() -> None:
@pytest.mark.anyio
@pytest.mark.parametrize(
"server_transport",
[
("tool_progress", "sse"),
("tool_progress", "streamable-http"),
],
indirect=True,
)
async def test_tool_progress(server_transport: str, server_url: str) -> None:
"""Test tool progress reporting.""" """Test tool progress reporting."""
transport = server_transport
collector = NotificationCollector() collector = NotificationCollector()
async def message_handler(message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception): async def message_handler(message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception):
@ -358,15 +192,7 @@ async def test_tool_progress(server_transport: str, server_url: str) -> None:
if isinstance(message, Exception): # pragma: no cover if isinstance(message, Exception): # pragma: no cover
raise message raise message
client_cm = create_client_for_transport(transport, server_url) async with Client(tool_progress.mcp, message_handler=message_handler) as client:
async with client_cm as (read_stream, write_stream):
async with ClientSession(read_stream, write_stream, message_handler=message_handler) as session:
# Test initialization
result = await session.initialize()
assert isinstance(result, InitializeResult)
assert result.server_info.name == "Progress Example"
# Test progress callback # Test progress callback
progress_updates = [] progress_updates = []
@ -375,7 +201,7 @@ async def test_tool_progress(server_transport: str, server_url: str) -> None:
# Call tool with progress # Call tool with progress
steps = 3 steps = 3
tool_result = await session.call_tool( tool_result = await client.call_tool(
"long_running_task", "long_running_task",
{"task_name": "Test Task", "steps": steps}, {"task_name": "Test Task", "steps": steps},
progress_callback=progress_callback, progress_callback=progress_callback,
@ -394,60 +220,24 @@ async def test_tool_progress(server_transport: str, server_url: str) -> None:
assert len(collector.log_messages) > 0 assert len(collector.log_messages) > 0
# Test sampling async def test_sampling() -> None:
@pytest.mark.anyio
@pytest.mark.parametrize(
"server_transport",
[
("sampling", "sse"),
("sampling", "streamable-http"),
],
indirect=True,
)
async def test_sampling(server_transport: str, server_url: str) -> None:
"""Test sampling (LLM interaction) functionality.""" """Test sampling (LLM interaction) functionality."""
transport = server_transport async with Client(sampling.mcp, sampling_callback=sampling_callback) as client:
client_cm = create_client_for_transport(transport, server_url) assert client.server_capabilities is not None
assert client.server_capabilities.tools is not None
async with client_cm as (read_stream, write_stream):
async with ClientSession(read_stream, write_stream, sampling_callback=sampling_callback) as session:
# Test initialization
result = await session.initialize()
assert isinstance(result, InitializeResult)
assert result.server_info.name == "Sampling Example"
assert result.capabilities.tools is not None
# Test sampling tool # Test sampling tool
sampling_result = await session.call_tool("generate_poem", {"topic": "nature"}) sampling_result = await client.call_tool("generate_poem", {"topic": "nature"})
assert len(sampling_result.content) == 1 assert len(sampling_result.content) == 1
assert isinstance(sampling_result.content[0], TextContent) assert isinstance(sampling_result.content[0], TextContent)
assert "This is a simulated LLM response" in sampling_result.content[0].text assert "This is a simulated LLM response" in sampling_result.content[0].text
# Test elicitation async def test_elicitation() -> None:
@pytest.mark.anyio
@pytest.mark.parametrize(
"server_transport",
[
("elicitation", "sse"),
("elicitation", "streamable-http"),
],
indirect=True,
)
async def test_elicitation(server_transport: str, server_url: str) -> None:
"""Test elicitation (user interaction) functionality.""" """Test elicitation (user interaction) functionality."""
transport = server_transport async with Client(elicitation.mcp, elicitation_callback=elicitation_callback) as client:
client_cm = create_client_for_transport(transport, server_url)
async with client_cm as (read_stream, write_stream):
async with ClientSession(read_stream, write_stream, elicitation_callback=elicitation_callback) as session:
# Test initialization
result = await session.initialize()
assert isinstance(result, InitializeResult)
assert result.server_info.name == "Elicitation Example"
# Test booking with unavailable date (triggers elicitation) # Test booking with unavailable date (triggers elicitation)
booking_result = await session.call_tool( booking_result = await client.call_tool(
"book_table", "book_table",
{ {
"date": "2024-12-25", # Unavailable date "date": "2024-12-25", # Unavailable date
@ -460,7 +250,7 @@ async def test_elicitation(server_transport: str, server_url: str) -> None:
assert "[SUCCESS] Booked for 2024-12-26" in booking_result.content[0].text assert "[SUCCESS] Booked for 2024-12-26" in booking_result.content[0].text
# Test booking with available date (no elicitation) # Test booking with available date (no elicitation)
booking_result = await session.call_tool( booking_result = await client.call_tool(
"book_table", "book_table",
{ {
"date": "2024-12-20", # Available date "date": "2024-12-20", # Available date
@ -473,19 +263,8 @@ async def test_elicitation(server_transport: str, server_url: str) -> None:
assert "[SUCCESS] Booked for 2024-12-20 at 20:00" in booking_result.content[0].text assert "[SUCCESS] Booked for 2024-12-20 at 20:00" in booking_result.content[0].text
# Test notifications async def test_notifications() -> None:
@pytest.mark.anyio
@pytest.mark.parametrize(
"server_transport",
[
("notifications", "sse"),
("notifications", "streamable-http"),
],
indirect=True,
)
async def test_notifications(server_transport: str, server_url: str) -> None:
"""Test notifications and logging functionality.""" """Test notifications and logging functionality."""
transport = server_transport
collector = NotificationCollector() collector = NotificationCollector()
async def message_handler(message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception): async def message_handler(message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception):
@ -493,17 +272,9 @@ async def test_notifications(server_transport: str, server_url: str) -> None:
if isinstance(message, Exception): # pragma: no cover if isinstance(message, Exception): # pragma: no cover
raise message raise message
client_cm = create_client_for_transport(transport, server_url) async with Client(notifications.mcp, message_handler=message_handler) as client:
async with client_cm as (read_stream, write_stream):
async with ClientSession(read_stream, write_stream, message_handler=message_handler) as session:
# Test initialization
result = await session.initialize()
assert isinstance(result, InitializeResult)
assert result.server_info.name == "Notifications Example"
# Call tool that generates notifications # Call tool that generates notifications
tool_result = await session.call_tool("process_data", {"data": "test_data"}) tool_result = await client.call_tool("process_data", {"data": "test_data"})
assert len(tool_result.content) == 1 assert len(tool_result.content) == 1
assert isinstance(tool_result.content[0], TextContent) assert isinstance(tool_result.content[0], TextContent)
assert "Processed: test_data" in tool_result.content[0].text assert "Processed: test_data" in tool_result.content[0].text
@ -520,32 +291,15 @@ async def test_notifications(server_transport: str, server_url: str) -> None:
assert len(collector.resource_notifications) > 0 assert len(collector.resource_notifications) > 0
# Test completion async def test_completion() -> None:
@pytest.mark.anyio
@pytest.mark.parametrize(
"server_transport",
[
("completion", "sse"),
("completion", "streamable-http"),
],
indirect=True,
)
async def test_completion(server_transport: str, server_url: str) -> None:
"""Test completion (autocomplete) functionality.""" """Test completion (autocomplete) functionality."""
transport = server_transport async with Client(completion.mcp) as client:
client_cm = create_client_for_transport(transport, server_url) assert client.server_capabilities is not None
assert client.server_capabilities.resources is not None
async with client_cm as (read_stream, write_stream): assert client.server_capabilities.prompts is not None
async with ClientSession(read_stream, write_stream) as session:
# Test initialization
result = await session.initialize()
assert isinstance(result, InitializeResult)
assert result.server_info.name == "Example"
assert result.capabilities.resources is not None
assert result.capabilities.prompts is not None
# Test resource completion # Test resource completion
completion_result = await session.complete( completion_result = await client.complete(
ref=ResourceTemplateReference(type="ref/resource", uri="github://repos/{owner}/{repo}"), ref=ResourceTemplateReference(type="ref/resource", uri="github://repos/{owner}/{repo}"),
argument={"name": "repo", "value": ""}, argument={"name": "repo", "value": ""},
context_arguments={"owner": "modelcontextprotocol"}, context_arguments={"owner": "modelcontextprotocol"},
@ -560,7 +314,7 @@ async def test_completion(server_transport: str, server_url: str) -> None:
assert "specification" in completion_result.completion.values assert "specification" in completion_result.completion.values
# Test prompt completion # Test prompt completion
completion_result = await session.complete( completion_result = await client.complete(
ref=PromptReference(type="ref/prompt", name="review_code"), ref=PromptReference(type="ref/prompt", name="review_code"),
argument={"name": "language", "value": "py"}, argument={"name": "language", "value": "py"},
) )
@ -572,65 +326,27 @@ async def test_completion(server_transport: str, server_url: str) -> None:
assert all(lang.startswith("py") for lang in completion_result.completion.values) assert all(lang.startswith("py") for lang in completion_result.completion.values)
# Test MCPServer quickstart example async def test_mcpserver_quickstart() -> None:
@pytest.mark.anyio
@pytest.mark.parametrize(
"server_transport",
[
("mcpserver_quickstart", "sse"),
("mcpserver_quickstart", "streamable-http"),
],
indirect=True,
)
async def test_mcpserver_quickstart(server_transport: str, server_url: str) -> None:
"""Test MCPServer quickstart example.""" """Test MCPServer quickstart example."""
transport = server_transport async with Client(mcpserver_quickstart.mcp) as client:
client_cm = create_client_for_transport(transport, server_url)
async with client_cm as (read_stream, write_stream):
async with ClientSession(read_stream, write_stream) as session:
# Test initialization
result = await session.initialize()
assert isinstance(result, InitializeResult)
assert result.server_info.name == "Demo"
# Test add tool # Test add tool
tool_result = await session.call_tool("add", {"a": 10, "b": 20}) tool_result = await client.call_tool("add", {"a": 10, "b": 20})
assert len(tool_result.content) == 1 assert len(tool_result.content) == 1
assert isinstance(tool_result.content[0], TextContent) assert isinstance(tool_result.content[0], TextContent)
assert tool_result.content[0].text == "30" assert tool_result.content[0].text == "30"
# Test greeting resource directly # Test greeting resource directly
resource_result = await session.read_resource("greeting://Alice") resource_result = await client.read_resource("greeting://Alice")
assert len(resource_result.contents) == 1 assert len(resource_result.contents) == 1
assert isinstance(resource_result.contents[0], TextResourceContents) assert isinstance(resource_result.contents[0], TextResourceContents)
assert resource_result.contents[0].text == "Hello, Alice!" assert resource_result.contents[0].text == "Hello, Alice!"
# Test structured output example async def test_structured_output() -> None:
@pytest.mark.anyio
@pytest.mark.parametrize(
"server_transport",
[
("structured_output", "sse"),
("structured_output", "streamable-http"),
],
indirect=True,
)
async def test_structured_output(server_transport: str, server_url: str) -> None:
"""Test structured output functionality.""" """Test structured output functionality."""
transport = server_transport async with Client(structured_output.mcp) as client:
client_cm = create_client_for_transport(transport, server_url)
async with client_cm as (read_stream, write_stream):
async with ClientSession(read_stream, write_stream) as session:
# Test initialization
result = await session.initialize()
assert isinstance(result, InitializeResult)
assert result.server_info.name == "Structured Output Example"
# Test get_weather tool # Test get_weather tool
weather_result = await session.call_tool("get_weather", {"city": "New York"}) weather_result = await client.call_tool("get_weather", {"city": "New York"})
assert len(weather_result.content) == 1 assert len(weather_result.content) == 1
assert isinstance(weather_result.content[0], TextContent) assert isinstance(weather_result.content[0], TextContent)