Add ergonomic Client class for testing MCP servers (#1870)
This commit is contained in:
Felix Weinberger
GitHub
parent
8adb5bdce8
commit
df039bf97c
7 changed files with 1298 additions and 526 deletions
|
|
@ -7,12 +7,11 @@ from typing import Any
|
|||
import pytest
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from mcp import types
|
||||
from mcp import Client, types
|
||||
from mcp.client.session import ClientSession, ElicitationFnT
|
||||
from mcp.server.fastmcp import Context, FastMCP
|
||||
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.types import ElicitRequestParams, ElicitResult, TextContent
|
||||
|
||||
|
||||
|
|
@ -47,12 +46,8 @@ async def call_tool_and_assert(
|
|||
text_contains: list[str] | None = None,
|
||||
):
|
||||
"""Helper to create session, call tool, and assert result."""
|
||||
async with create_connected_server_and_client_session(
|
||||
mcp._mcp_server, elicitation_callback=elicitation_callback
|
||||
) as client_session:
|
||||
await client_session.initialize()
|
||||
|
||||
result = await client_session.call_tool(tool_name, args)
|
||||
async with Client(mcp, elicitation_callback=elicitation_callback) as client:
|
||||
result = await client.call_tool(tool_name, args)
|
||||
assert len(result.content) == 1
|
||||
assert isinstance(result.content[0], TextContent)
|
||||
|
||||
|
|
@ -134,14 +129,10 @@ async def test_elicitation_schema_validation():
|
|||
): # pragma: no cover
|
||||
return ElicitResult(action="accept", content={})
|
||||
|
||||
async with create_connected_server_and_client_session(
|
||||
mcp._mcp_server, elicitation_callback=elicitation_callback
|
||||
) as client_session:
|
||||
await client_session.initialize()
|
||||
|
||||
async with Client(mcp, elicitation_callback=elicitation_callback) as client:
|
||||
# Test both invalid schemas
|
||||
for tool_name, field_name in [("invalid_list", "numbers"), ("nested_model", "nested")]:
|
||||
result = await client_session.call_tool(tool_name, {})
|
||||
result = await client.call_tool(tool_name, {})
|
||||
assert len(result.content) == 1
|
||||
assert isinstance(result.content[0], TextContent)
|
||||
assert "Validation failed as expected" in result.content[0].text
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from pydantic import BaseModel
|
|||
from starlette.applications import Starlette
|
||||
from starlette.routing import Mount, Route
|
||||
|
||||
from mcp.client import Client
|
||||
from mcp.server.fastmcp import Context, FastMCP
|
||||
from mcp.server.fastmcp.exceptions import ToolError
|
||||
from mcp.server.fastmcp.prompts.base import Message, UserMessage
|
||||
|
|
@ -16,7 +17,6 @@ from mcp.server.fastmcp.utilities.types import Audio, Image
|
|||
from mcp.server.session import ServerSession
|
||||
from mcp.server.transport_security import TransportSecuritySettings
|
||||
from mcp.shared.exceptions import McpError
|
||||
from mcp.shared.memory import create_connected_server_and_client_session as client_session
|
||||
from mcp.types import (
|
||||
AudioContent,
|
||||
BlobResourceContents,
|
||||
|
|
@ -77,7 +77,7 @@ class TestServer:
|
|||
def hello_world(name: str = "世界") -> str:
|
||||
return f"¡Hola, {name}! 👋"
|
||||
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
tools = await client.list_tools()
|
||||
assert len(tools.tools) == 1
|
||||
tool = tools.tools[0]
|
||||
|
|
@ -230,7 +230,7 @@ class TestServerTools:
|
|||
async def test_list_tools(self):
|
||||
mcp = FastMCP()
|
||||
mcp.add_tool(tool_fn)
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
tools = await client.list_tools()
|
||||
assert len(tools.tools) == 1
|
||||
|
||||
|
|
@ -238,7 +238,7 @@ class TestServerTools:
|
|||
async def test_call_tool(self):
|
||||
mcp = FastMCP()
|
||||
mcp.add_tool(tool_fn)
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
result = await client.call_tool("my_tool", {"arg1": "value"})
|
||||
assert not hasattr(result, "error")
|
||||
assert len(result.content) > 0
|
||||
|
|
@ -247,7 +247,7 @@ class TestServerTools:
|
|||
async def test_tool_exception_handling(self):
|
||||
mcp = FastMCP()
|
||||
mcp.add_tool(error_tool_fn)
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
result = await client.call_tool("error_tool_fn", {})
|
||||
assert len(result.content) == 1
|
||||
content = result.content[0]
|
||||
|
|
@ -259,7 +259,7 @@ class TestServerTools:
|
|||
async def test_tool_error_handling(self):
|
||||
mcp = FastMCP()
|
||||
mcp.add_tool(error_tool_fn)
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
result = await client.call_tool("error_tool_fn", {})
|
||||
assert len(result.content) == 1
|
||||
content = result.content[0]
|
||||
|
|
@ -272,7 +272,7 @@ class TestServerTools:
|
|||
"""Test that exception details are properly formatted in the response"""
|
||||
mcp = FastMCP()
|
||||
mcp.add_tool(error_tool_fn)
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
result = await client.call_tool("error_tool_fn", {})
|
||||
content = result.content[0]
|
||||
assert isinstance(content, TextContent)
|
||||
|
|
@ -284,7 +284,7 @@ class TestServerTools:
|
|||
async def test_tool_return_value_conversion(self):
|
||||
mcp = FastMCP()
|
||||
mcp.add_tool(tool_fn)
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
result = await client.call_tool("tool_fn", {"x": 1, "y": 2})
|
||||
assert len(result.content) == 1
|
||||
content = result.content[0]
|
||||
|
|
@ -302,7 +302,7 @@ class TestServerTools:
|
|||
|
||||
mcp = FastMCP()
|
||||
mcp.add_tool(image_tool_fn)
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
result = await client.call_tool("image_tool_fn", {"path": str(image_path)})
|
||||
assert len(result.content) == 1
|
||||
content = result.content[0]
|
||||
|
|
@ -323,7 +323,7 @@ class TestServerTools:
|
|||
|
||||
mcp = FastMCP()
|
||||
mcp.add_tool(audio_tool_fn)
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
result = await client.call_tool("audio_tool_fn", {"path": str(audio_path)})
|
||||
assert len(result.content) == 1
|
||||
content = result.content[0]
|
||||
|
|
@ -358,7 +358,7 @@ class TestServerTools:
|
|||
audio_path = tmp_path / filename
|
||||
audio_path.write_bytes(b"fake audio data")
|
||||
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
result = await client.call_tool("audio_tool_fn", {"path": str(audio_path)})
|
||||
assert len(result.content) == 1
|
||||
content = result.content[0]
|
||||
|
|
@ -373,7 +373,7 @@ class TestServerTools:
|
|||
async def test_tool_mixed_content(self):
|
||||
mcp = FastMCP()
|
||||
mcp.add_tool(mixed_content_tool_fn)
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
result = await client.call_tool("mixed_content_tool_fn", {})
|
||||
assert len(result.content) == 3
|
||||
content1, content2, content3 = result.content
|
||||
|
|
@ -425,7 +425,7 @@ class TestServerTools:
|
|||
|
||||
mcp = FastMCP()
|
||||
mcp.add_tool(mixed_list_fn) # type: ignore
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
result = await client.call_tool("mixed_list_fn", {})
|
||||
assert len(result.content) == 5
|
||||
# Check text conversion
|
||||
|
|
@ -469,7 +469,7 @@ class TestServerTools:
|
|||
mcp = FastMCP()
|
||||
mcp.add_tool(get_user)
|
||||
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
# Check that the tool has outputSchema
|
||||
tools = await client.list_tools()
|
||||
tool = next(t for t in tools.tools if t.name == "get_user")
|
||||
|
|
@ -499,7 +499,7 @@ class TestServerTools:
|
|||
mcp = FastMCP()
|
||||
mcp.add_tool(calculate_sum)
|
||||
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
# Check that the tool has outputSchema
|
||||
tools = await client.list_tools()
|
||||
tool = next(t for t in tools.tools if t.name == "calculate_sum")
|
||||
|
|
@ -526,7 +526,7 @@ class TestServerTools:
|
|||
mcp = FastMCP()
|
||||
mcp.add_tool(get_numbers)
|
||||
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
result = await client.call_tool("get_numbers", {})
|
||||
assert result.is_error is False
|
||||
assert result.structured_content is not None
|
||||
|
|
@ -542,7 +542,7 @@ class TestServerTools:
|
|||
mcp = FastMCP()
|
||||
mcp.add_tool(get_numbers)
|
||||
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
result = await client.call_tool("get_numbers", {})
|
||||
assert result.is_error is True
|
||||
assert result.structured_content is None
|
||||
|
|
@ -566,7 +566,7 @@ class TestServerTools:
|
|||
mcp = FastMCP()
|
||||
mcp.add_tool(get_metadata)
|
||||
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
# Check schema
|
||||
tools = await client.list_tools()
|
||||
tool = next(t for t in tools.tools if t.name == "get_metadata")
|
||||
|
|
@ -602,7 +602,7 @@ class TestServerTools:
|
|||
mcp = FastMCP()
|
||||
mcp.add_tool(get_settings)
|
||||
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
# Check schema
|
||||
tools = await client.list_tools()
|
||||
tool = next(t for t in tools.tools if t.name == "get_settings")
|
||||
|
|
@ -646,7 +646,7 @@ class TestServerTools:
|
|||
mcp.add_tool(error_tool_fn)
|
||||
|
||||
# Verify both tools exist
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
tools = await client.list_tools()
|
||||
assert len(tools.tools) == 2
|
||||
tool_names = [t.name for t in tools.tools]
|
||||
|
|
@ -657,7 +657,7 @@ class TestServerTools:
|
|||
mcp.remove_tool("tool_fn")
|
||||
|
||||
# Verify only one tool remains
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
tools = await client.list_tools()
|
||||
assert len(tools.tools) == 1
|
||||
assert tools.tools[0].name == "error_tool_fn"
|
||||
|
|
@ -669,7 +669,7 @@ class TestServerTools:
|
|||
mcp.add_tool(tool_fn)
|
||||
|
||||
# Verify tool works before removal
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
result = await client.call_tool("tool_fn", {"x": 1, "y": 2})
|
||||
assert not result.is_error
|
||||
content = result.content[0]
|
||||
|
|
@ -680,7 +680,7 @@ class TestServerTools:
|
|||
mcp.remove_tool("tool_fn")
|
||||
|
||||
# Verify calling removed tool returns an error
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
result = await client.call_tool("tool_fn", {"x": 1, "y": 2})
|
||||
assert result.is_error
|
||||
content = result.content[0]
|
||||
|
|
@ -699,8 +699,12 @@ class TestServerResources:
|
|||
resource = FunctionResource(uri="resource://test", name="test", fn=get_text)
|
||||
mcp.add_resource(resource)
|
||||
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
result = await client.read_resource("resource://test")
|
||||
|
||||
async with Client(mcp) as client:
|
||||
result = await client.read_resource("resource://test")
|
||||
|
||||
assert isinstance(result.contents[0], TextResourceContents)
|
||||
assert result.contents[0].text == "Hello, world!"
|
||||
|
||||
|
|
@ -719,8 +723,12 @@ class TestServerResources:
|
|||
)
|
||||
mcp.add_resource(resource)
|
||||
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
result = await client.read_resource("resource://binary")
|
||||
|
||||
async with Client(mcp) as client:
|
||||
result = await client.read_resource("resource://binary")
|
||||
|
||||
assert isinstance(result.contents[0], BlobResourceContents)
|
||||
assert result.contents[0].blob == base64.b64encode(b"Binary data").decode()
|
||||
|
||||
|
|
@ -735,8 +743,12 @@ class TestServerResources:
|
|||
resource = FileResource(uri="file://test.txt", name="test.txt", path=text_file)
|
||||
mcp.add_resource(resource)
|
||||
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
result = await client.read_resource("file://test.txt")
|
||||
|
||||
async with Client(mcp) as client:
|
||||
result = await client.read_resource("file://test.txt")
|
||||
|
||||
assert isinstance(result.contents[0], TextResourceContents)
|
||||
assert result.contents[0].text == "Hello from file!"
|
||||
|
||||
|
|
@ -756,8 +768,12 @@ class TestServerResources:
|
|||
)
|
||||
mcp.add_resource(resource)
|
||||
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
result = await client.read_resource("file://test.bin")
|
||||
|
||||
async with Client(mcp) as client:
|
||||
result = await client.read_resource("file://test.bin")
|
||||
|
||||
assert isinstance(result.contents[0], BlobResourceContents)
|
||||
assert result.contents[0].blob == base64.b64encode(b"Binary file data").decode()
|
||||
|
||||
|
|
@ -770,7 +786,7 @@ class TestServerResources:
|
|||
"""get_data returns a string"""
|
||||
return "Hello, world!"
|
||||
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
resources = await client.list_resources()
|
||||
assert len(resources.resources) == 1
|
||||
resource = resources.resources[0]
|
||||
|
|
@ -822,8 +838,12 @@ class TestServerResourceTemplates:
|
|||
def get_data(name: str) -> str:
|
||||
return f"Data for {name}"
|
||||
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
result = await client.read_resource("resource://test/data")
|
||||
|
||||
async with Client(mcp) as client:
|
||||
result = await client.read_resource("resource://test/data")
|
||||
|
||||
assert isinstance(result.contents[0], TextResourceContents)
|
||||
assert result.contents[0].text == "Data for test"
|
||||
|
||||
|
|
@ -847,8 +867,12 @@ class TestServerResourceTemplates:
|
|||
def get_data(org: str, repo: str) -> str:
|
||||
return f"Data for {org}/{repo}"
|
||||
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
result = await client.read_resource("resource://cursor/fastmcp/data")
|
||||
|
||||
async with Client(mcp) as client:
|
||||
result = await client.read_resource("resource://cursor/fastmcp/data")
|
||||
|
||||
assert isinstance(result.contents[0], TextResourceContents)
|
||||
assert result.contents[0].text == "Data for cursor/fastmcp"
|
||||
|
||||
|
|
@ -870,8 +894,12 @@ class TestServerResourceTemplates:
|
|||
def get_static_data() -> str:
|
||||
return "Static data"
|
||||
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
result = await client.read_resource("resource://static")
|
||||
|
||||
async with Client(mcp) as client:
|
||||
result = await client.read_resource("resource://static")
|
||||
|
||||
assert isinstance(result.contents[0], TextResourceContents)
|
||||
assert result.contents[0].text == "Static data"
|
||||
|
||||
|
|
@ -910,8 +938,12 @@ class TestServerResourceTemplates:
|
|||
assert hasattr(template, "mime_type")
|
||||
assert template.mime_type == "text/csv"
|
||||
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
result = await client.read_resource("resource://bob/csv")
|
||||
|
||||
async with Client(mcp) as client:
|
||||
result = await client.read_resource("resource://bob/csv")
|
||||
|
||||
assert isinstance(result.contents[0], TextResourceContents)
|
||||
assert result.contents[0].text == "csv for bob"
|
||||
|
||||
|
|
@ -972,7 +1004,10 @@ class TestServerResourceMetadata:
|
|||
def get_data() -> str:
|
||||
return "test data"
|
||||
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
result = await client.read_resource("resource://data")
|
||||
|
||||
async with Client(mcp) as client:
|
||||
result = await client.read_resource("resource://data")
|
||||
|
||||
# Verify content and metadata in protocol response
|
||||
|
|
@ -1008,7 +1043,7 @@ class TestContextInjection:
|
|||
return f"Request {ctx.request_id}: {x}"
|
||||
|
||||
mcp.add_tool(tool_with_context)
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
result = await client.call_tool("tool_with_context", {"x": 42})
|
||||
assert len(result.content) == 1
|
||||
content = result.content[0]
|
||||
|
|
@ -1026,7 +1061,7 @@ class TestContextInjection:
|
|||
return f"Async request {ctx.request_id}: {x}"
|
||||
|
||||
mcp.add_tool(async_tool)
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
result = await client.call_tool("async_tool", {"x": 42})
|
||||
assert len(result.content) == 1
|
||||
content = result.content[0]
|
||||
|
|
@ -1049,7 +1084,7 @@ class TestContextInjection:
|
|||
mcp.add_tool(logging_tool)
|
||||
|
||||
with patch("mcp.server.session.ServerSession.send_log_message") as mock_log:
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
result = await client.call_tool("logging_tool", {"msg": "test"})
|
||||
assert len(result.content) == 1
|
||||
content = result.content[0]
|
||||
|
|
@ -1091,7 +1126,7 @@ class TestContextInjection:
|
|||
return x * 2
|
||||
|
||||
mcp.add_tool(no_context)
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
result = await client.call_tool("no_context", {"x": 21})
|
||||
assert len(result.content) == 1
|
||||
content = result.content[0]
|
||||
|
|
@ -1115,7 +1150,7 @@ class TestContextInjection:
|
|||
r = r_list[0]
|
||||
return f"Read resource: {r.content} with mime type {r.mime_type}"
|
||||
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
result = await client.call_tool("tool_with_resource", {})
|
||||
assert len(result.content) == 1
|
||||
content = result.content[0]
|
||||
|
|
@ -1141,8 +1176,13 @@ class TestContextInjection:
|
|||
assert template.context_kwarg == "ctx"
|
||||
|
||||
# Test via client
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
|
||||
async with Client(mcp) as client:
|
||||
result = await client.read_resource("resource://context/test")
|
||||
|
||||
async with Client(mcp) as client:
|
||||
result = await client.read_resource("resource://context/test")
|
||||
|
||||
assert len(result.contents) == 1
|
||||
content = result.contents[0]
|
||||
assert isinstance(content, TextResourceContents)
|
||||
|
|
@ -1166,8 +1206,13 @@ class TestContextInjection:
|
|||
assert template.context_kwarg is None
|
||||
|
||||
# Test via client
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
|
||||
async with Client(mcp) as client:
|
||||
result = await client.read_resource("resource://nocontext/test")
|
||||
|
||||
async with Client(mcp) as client:
|
||||
result = await client.read_resource("resource://nocontext/test")
|
||||
|
||||
assert len(result.contents) == 1
|
||||
content = result.contents[0]
|
||||
assert isinstance(content, TextResourceContents)
|
||||
|
|
@ -1191,8 +1236,13 @@ class TestContextInjection:
|
|||
assert template.context_kwarg == "my_ctx"
|
||||
|
||||
# Test via client
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
|
||||
async with Client(mcp) as client:
|
||||
result = await client.read_resource("resource://custom/123")
|
||||
|
||||
async with Client(mcp) as client:
|
||||
result = await client.read_resource("resource://custom/123")
|
||||
|
||||
assert len(result.contents) == 1
|
||||
content = result.contents[0]
|
||||
assert isinstance(content, TextResourceContents)
|
||||
|
|
@ -1214,7 +1264,7 @@ class TestContextInjection:
|
|||
assert len(prompts) == 1
|
||||
|
||||
# Test via client
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
# Try calling without passing ctx explicitly
|
||||
result = await client.get_prompt("prompt_with_ctx", {"text": "test"})
|
||||
# If this succeeds, check if context was injected
|
||||
|
|
@ -1234,7 +1284,7 @@ class TestContextInjection:
|
|||
return f"Prompt '{text}' works"
|
||||
|
||||
# Test via client
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
result = await client.get_prompt("prompt_no_ctx", {"text": "test"})
|
||||
assert len(result.messages) == 1
|
||||
message = result.messages[0]
|
||||
|
|
@ -1313,7 +1363,7 @@ class TestServerPrompts:
|
|||
def fn(name: str, optional: str = "default") -> str: # pragma: no cover
|
||||
return f"Hello, {name}!"
|
||||
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
result = await client.list_prompts()
|
||||
assert result.prompts is not None
|
||||
assert len(result.prompts) == 1
|
||||
|
|
@ -1335,7 +1385,7 @@ class TestServerPrompts:
|
|||
def fn(name: str) -> str:
|
||||
return f"Hello, {name}!"
|
||||
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
result = await client.get_prompt("fn", {"name": "World"})
|
||||
assert len(result.messages) == 1
|
||||
message = result.messages[0]
|
||||
|
|
@ -1353,7 +1403,7 @@ class TestServerPrompts:
|
|||
def fn(name: str) -> str:
|
||||
return f"Hello, {name}!"
|
||||
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
result = await client.get_prompt("fn", {"name": "World"})
|
||||
assert result.description == "Test prompt description"
|
||||
|
||||
|
|
@ -1366,7 +1416,7 @@ class TestServerPrompts:
|
|||
def fn(name: str) -> str:
|
||||
return f"Hello, {name}!"
|
||||
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
result = await client.get_prompt("fn", {"name": "World"})
|
||||
assert result.description == ""
|
||||
|
||||
|
|
@ -1380,7 +1430,7 @@ class TestServerPrompts:
|
|||
"""This is the function docstring."""
|
||||
return f"Hello, {name}!"
|
||||
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
result = await client.get_prompt("fn", {"name": "World"})
|
||||
assert result.description == "This is the function docstring."
|
||||
|
||||
|
|
@ -1402,7 +1452,7 @@ class TestServerPrompts:
|
|||
)
|
||||
)
|
||||
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
result = await client.get_prompt("fn")
|
||||
assert len(result.messages) == 1
|
||||
message = result.messages[0]
|
||||
|
|
@ -1418,7 +1468,7 @@ class TestServerPrompts:
|
|||
async def test_get_unknown_prompt(self):
|
||||
"""Test error when getting unknown prompt."""
|
||||
mcp = FastMCP()
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
with pytest.raises(McpError, match="Unknown prompt"):
|
||||
await client.get_prompt("unknown")
|
||||
|
||||
|
|
@ -1431,7 +1481,7 @@ class TestServerPrompts:
|
|||
def prompt_fn(name: str) -> str: # pragma: no cover
|
||||
return f"Hello, {name}!"
|
||||
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
async with Client(mcp) as client:
|
||||
with pytest.raises(McpError, match="Missing required arguments"):
|
||||
await client.get_prompt("prompt_fn")
|
||||
|
||||
|
|
|
|||
|
|
@ -2,9 +2,9 @@
|
|||
|
||||
import pytest
|
||||
|
||||
from mcp import Client
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
from mcp.server.fastmcp.resources import FunctionResource
|
||||
from mcp.shared.memory import create_connected_server_and_client_session
|
||||
from mcp.shared.metadata_utils import get_display_name
|
||||
from mcp.types import Prompt, Resource, ResourceTemplate, Tool, ToolAnnotations
|
||||
|
||||
|
|
@ -24,8 +24,9 @@ async def test_server_name_title_description_version():
|
|||
assert mcp.version == "1.0"
|
||||
|
||||
# Start server and connect client
|
||||
async with create_connected_server_and_client_session(mcp._mcp_server) as client:
|
||||
init_result = await client.initialize()
|
||||
async with Client(mcp) as client:
|
||||
# Access initialization result from session
|
||||
init_result = await client.session.initialize()
|
||||
assert init_result.server_info.name == "TestServer"
|
||||
assert init_result.server_info.title == "Test Server Title"
|
||||
assert init_result.server_info.description == "This is a test server description."
|
||||
|
|
@ -60,9 +61,7 @@ async def test_tool_title_precedence():
|
|||
return message
|
||||
|
||||
# Start server and connect client
|
||||
async with create_connected_server_and_client_session(mcp._mcp_server) as client:
|
||||
await client.initialize()
|
||||
|
||||
async with Client(mcp) as client:
|
||||
# List tools
|
||||
tools_result = await client.list_tools()
|
||||
tools = {tool.name: tool for tool in tools_result.tools}
|
||||
|
|
@ -104,9 +103,7 @@ async def test_prompt_title():
|
|||
return f"Tell me about {topic}"
|
||||
|
||||
# Start server and connect client
|
||||
async with create_connected_server_and_client_session(mcp._mcp_server) as client:
|
||||
await client.initialize()
|
||||
|
||||
async with Client(mcp) as client:
|
||||
# List prompts
|
||||
prompts_result = await client.list_prompts()
|
||||
prompts = {prompt.name: prompt for prompt in prompts_result.prompts}
|
||||
|
|
@ -164,9 +161,7 @@ async def test_resource_title():
|
|||
return f"Data for {id}"
|
||||
|
||||
# Start server and connect client
|
||||
async with create_connected_server_and_client_session(mcp._mcp_server) as client:
|
||||
await client.initialize()
|
||||
|
||||
async with Client(mcp) as client:
|
||||
# List resources
|
||||
resources_result = await client.list_resources()
|
||||
resources = {str(res.uri): res for res in resources_result.resources}
|
||||
|
|
|
|||
|
|
@ -4,13 +4,12 @@ import anyio
|
|||
import pytest
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from mcp import types
|
||||
from mcp import Client, types
|
||||
from mcp.client.session import ClientSession
|
||||
from mcp.server.elicitation import CancelledElicitation, DeclinedElicitation, elicit_url
|
||||
from mcp.server.fastmcp import Context, FastMCP
|
||||
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.types import ElicitRequestParams, ElicitResult, TextContent
|
||||
|
||||
|
||||
|
|
@ -37,12 +36,8 @@ async def test_url_elicitation_accept():
|
|||
assert params.message == "Please provide your API key to continue."
|
||||
return ElicitResult(action="accept")
|
||||
|
||||
async with create_connected_server_and_client_session(
|
||||
mcp._mcp_server, elicitation_callback=elicitation_callback
|
||||
) as client_session:
|
||||
await client_session.initialize()
|
||||
|
||||
result = await client_session.call_tool("request_api_key", {})
|
||||
async with Client(mcp, elicitation_callback=elicitation_callback) as client:
|
||||
result = await client.call_tool("request_api_key", {})
|
||||
assert len(result.content) == 1
|
||||
assert isinstance(result.content[0], TextContent)
|
||||
assert result.content[0].text == "User accept"
|
||||
|
|
@ -67,12 +62,8 @@ async def test_url_elicitation_decline():
|
|||
assert params.mode == "url"
|
||||
return ElicitResult(action="decline")
|
||||
|
||||
async with create_connected_server_and_client_session(
|
||||
mcp._mcp_server, elicitation_callback=elicitation_callback
|
||||
) as client_session:
|
||||
await client_session.initialize()
|
||||
|
||||
result = await client_session.call_tool("oauth_flow", {})
|
||||
async with Client(mcp, elicitation_callback=elicitation_callback) as client:
|
||||
result = await client.call_tool("oauth_flow", {})
|
||||
assert len(result.content) == 1
|
||||
assert isinstance(result.content[0], TextContent)
|
||||
assert result.content[0].text == "User decline authorization"
|
||||
|
|
@ -97,12 +88,8 @@ async def test_url_elicitation_cancel():
|
|||
assert params.mode == "url"
|
||||
return ElicitResult(action="cancel")
|
||||
|
||||
async with create_connected_server_and_client_session(
|
||||
mcp._mcp_server, elicitation_callback=elicitation_callback
|
||||
) as client_session:
|
||||
await client_session.initialize()
|
||||
|
||||
result = await client_session.call_tool("payment_flow", {})
|
||||
async with Client(mcp, elicitation_callback=elicitation_callback) as client:
|
||||
result = await client.call_tool("payment_flow", {})
|
||||
assert len(result.content) == 1
|
||||
assert isinstance(result.content[0], TextContent)
|
||||
assert result.content[0].text == "User cancel payment"
|
||||
|
|
@ -127,12 +114,8 @@ async def test_url_elicitation_helper_function():
|
|||
async def elicitation_callback(context: RequestContext[ClientSession, None], params: ElicitRequestParams):
|
||||
return ElicitResult(action="accept")
|
||||
|
||||
async with create_connected_server_and_client_session(
|
||||
mcp._mcp_server, elicitation_callback=elicitation_callback
|
||||
) as client_session:
|
||||
await client_session.initialize()
|
||||
|
||||
result = await client_session.call_tool("setup_credentials", {})
|
||||
async with Client(mcp, elicitation_callback=elicitation_callback) as client:
|
||||
result = await client.call_tool("setup_credentials", {})
|
||||
assert len(result.content) == 1
|
||||
assert isinstance(result.content[0], TextContent)
|
||||
assert result.content[0].text == "AcceptedUrlElicitation"
|
||||
|
|
@ -165,12 +148,8 @@ async def test_url_no_content_in_response():
|
|||
# Return without content - this is correct for URL mode
|
||||
return ElicitResult(action="accept")
|
||||
|
||||
async with create_connected_server_and_client_session(
|
||||
mcp._mcp_server, elicitation_callback=elicitation_callback
|
||||
) as client_session:
|
||||
await client_session.initialize()
|
||||
|
||||
result = await client_session.call_tool("check_url_response", {})
|
||||
async with Client(mcp, elicitation_callback=elicitation_callback) as client:
|
||||
result = await client.call_tool("check_url_response", {})
|
||||
assert len(result.content) == 1
|
||||
assert isinstance(result.content[0], TextContent)
|
||||
assert "Content: None" in result.content[0].text
|
||||
|
|
@ -200,12 +179,8 @@ async def test_form_mode_still_works():
|
|||
assert params.requested_schema is not None
|
||||
return ElicitResult(action="accept", content={"name": "Alice"})
|
||||
|
||||
async with create_connected_server_and_client_session(
|
||||
mcp._mcp_server, elicitation_callback=elicitation_callback
|
||||
) as client_session:
|
||||
await client_session.initialize()
|
||||
|
||||
result = await client_session.call_tool("ask_name", {})
|
||||
async with Client(mcp, elicitation_callback=elicitation_callback) as client:
|
||||
result = await client.call_tool("ask_name", {})
|
||||
assert len(result.content) == 1
|
||||
assert isinstance(result.content[0], TextContent)
|
||||
assert result.content[0].text == "Hello, Alice!"
|
||||
|
|
@ -235,12 +210,8 @@ async def test_elicit_complete_notification():
|
|||
async def elicitation_callback(context: RequestContext[ClientSession, None], params: ElicitRequestParams):
|
||||
return ElicitResult(action="accept") # pragma: no cover
|
||||
|
||||
async with create_connected_server_and_client_session(
|
||||
mcp._mcp_server, elicitation_callback=elicitation_callback
|
||||
) as client_session:
|
||||
await client_session.initialize()
|
||||
|
||||
result = await client_session.call_tool("trigger_elicitation", {})
|
||||
async with Client(mcp, elicitation_callback=elicitation_callback) as client:
|
||||
result = await client.call_tool("trigger_elicitation", {})
|
||||
assert len(result.content) == 1
|
||||
assert isinstance(result.content[0], TextContent)
|
||||
assert result.content[0].text == "Elicitation completed"
|
||||
|
|
@ -296,12 +267,8 @@ async def test_elicit_url_typed_results():
|
|||
async def decline_callback(context: RequestContext[ClientSession, None], params: ElicitRequestParams):
|
||||
return ElicitResult(action="decline")
|
||||
|
||||
async with create_connected_server_and_client_session(
|
||||
mcp._mcp_server, elicitation_callback=decline_callback
|
||||
) as client_session:
|
||||
await client_session.initialize()
|
||||
|
||||
result = await client_session.call_tool("test_decline", {})
|
||||
async with Client(mcp, elicitation_callback=decline_callback) as client:
|
||||
result = await client.call_tool("test_decline", {})
|
||||
assert len(result.content) == 1
|
||||
assert isinstance(result.content[0], TextContent)
|
||||
assert result.content[0].text == "Declined"
|
||||
|
|
@ -310,12 +277,8 @@ async def test_elicit_url_typed_results():
|
|||
async def cancel_callback(context: RequestContext[ClientSession, None], params: ElicitRequestParams):
|
||||
return ElicitResult(action="cancel")
|
||||
|
||||
async with create_connected_server_and_client_session(
|
||||
mcp._mcp_server, elicitation_callback=cancel_callback
|
||||
) as client_session:
|
||||
await client_session.initialize()
|
||||
|
||||
result = await client_session.call_tool("test_cancel", {})
|
||||
async with Client(mcp, elicitation_callback=cancel_callback) as client:
|
||||
result = await client.call_tool("test_cancel", {})
|
||||
assert len(result.content) == 1
|
||||
assert isinstance(result.content[0], TextContent)
|
||||
assert result.content[0].text == "Cancelled"
|
||||
|
|
@ -347,12 +310,8 @@ async def test_deprecated_elicit_method():
|
|||
assert params.requested_schema is not None
|
||||
return ElicitResult(action="accept", content={"email": "test@example.com"})
|
||||
|
||||
async with create_connected_server_and_client_session(
|
||||
mcp._mcp_server, elicitation_callback=elicitation_callback
|
||||
) as client_session:
|
||||
await client_session.initialize()
|
||||
|
||||
result = await client_session.call_tool("use_deprecated_elicit", {})
|
||||
async with Client(mcp, elicitation_callback=elicitation_callback) as client:
|
||||
result = await client.call_tool("use_deprecated_elicit", {})
|
||||
assert len(result.content) == 1
|
||||
assert isinstance(result.content[0], TextContent)
|
||||
assert result.content[0].text == "Email: test@example.com"
|
||||
|
|
@ -378,10 +337,7 @@ async def test_ctx_elicit_url_convenience_method():
|
|||
assert params.elicitation_id == "ctx-test-001"
|
||||
return ElicitResult(action="accept")
|
||||
|
||||
async with create_connected_server_and_client_session(
|
||||
mcp._mcp_server, elicitation_callback=elicitation_callback
|
||||
) as client_session:
|
||||
await client_session.initialize()
|
||||
result = await client_session.call_tool("direct_elicit_url", {})
|
||||
async with Client(mcp, elicitation_callback=elicitation_callback) as client:
|
||||
result = await client.call_tool("direct_elicit_url", {})
|
||||
assert isinstance(result.content[0], TextContent)
|
||||
assert result.content[0].text == "Result: accept"
|
||||
|
|
|
|||
|
|
@ -2,11 +2,10 @@
|
|||
|
||||
import pytest
|
||||
|
||||
from mcp import types
|
||||
from mcp import Client, types
|
||||
from mcp.server.fastmcp import Context, FastMCP
|
||||
from mcp.server.session import ServerSession
|
||||
from mcp.shared.exceptions import McpError, UrlElicitationRequiredError
|
||||
from mcp.shared.memory import create_connected_server_and_client_session
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
|
|
@ -28,12 +27,10 @@ async def test_url_elicitation_error_thrown_from_tool():
|
|||
]
|
||||
)
|
||||
|
||||
async with create_connected_server_and_client_session(mcp._mcp_server) as client_session:
|
||||
await client_session.initialize()
|
||||
|
||||
async with Client(mcp) as client:
|
||||
# Call the tool - it should raise McpError with URL_ELICITATION_REQUIRED code
|
||||
with pytest.raises(McpError) as exc_info:
|
||||
await client_session.call_tool("connect_service", {"service_name": "github"})
|
||||
await client.call_tool("connect_service", {"service_name": "github"})
|
||||
|
||||
# Verify the error details
|
||||
error = exc_info.value.error
|
||||
|
|
@ -74,12 +71,10 @@ async def test_url_elicitation_error_from_error():
|
|||
]
|
||||
)
|
||||
|
||||
async with create_connected_server_and_client_session(mcp._mcp_server) as client_session:
|
||||
await client_session.initialize()
|
||||
|
||||
async with Client(mcp) as client:
|
||||
# Call the tool and catch the error
|
||||
with pytest.raises(McpError) as exc_info:
|
||||
await client_session.call_tool("multi_auth", {})
|
||||
await client.call_tool("multi_auth", {})
|
||||
|
||||
# Reconstruct the typed error
|
||||
mcp_error = exc_info.value
|
||||
|
|
@ -102,11 +97,9 @@ async def test_normal_exceptions_still_return_error_result():
|
|||
async def failing_tool(ctx: Context[ServerSession, None]) -> str:
|
||||
raise ValueError("Something went wrong")
|
||||
|
||||
async with create_connected_server_and_client_session(mcp._mcp_server) as client_session:
|
||||
await client_session.initialize()
|
||||
|
||||
async with Client(mcp) as client:
|
||||
# Normal exceptions should be returned as error results, not McpError
|
||||
result = await client_session.call_tool("failing_tool", {})
|
||||
result = await client.call_tool("failing_tool", {})
|
||||
assert result.is_error is True
|
||||
assert len(result.content) == 1
|
||||
assert isinstance(result.content[0], types.TextContent)
|
||||
|
|
|
|||
|
|
@ -6,9 +6,10 @@ 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_connected_server_and_client_session
|
||||
from mcp.types import (
|
||||
CallToolRequest,
|
||||
CallToolRequestParams,
|
||||
|
|
@ -54,57 +55,61 @@ async def test_server_remains_functional_after_cancel():
|
|||
return [types.TextContent(type="text", text=f"Call number: {call_count}")]
|
||||
raise ValueError(f"Unknown tool: {name}") # pragma: no cover
|
||||
|
||||
async with create_connected_server_and_client_session(server) as client:
|
||||
# First request (will be cancelled)
|
||||
async def first_request():
|
||||
try:
|
||||
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
|
||||
transport = InMemoryTransport(server)
|
||||
async with transport.connect() as (read_stream, write_stream):
|
||||
async with ClientSession(read_stream, write_stream) as client:
|
||||
await client.initialize()
|
||||
|
||||
# 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",
|
||||
# First request (will be cancelled)
|
||||
async def first_request():
|
||||
try:
|
||||
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",
|
||||
),
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
# Second request (should work normally)
|
||||
result = await client.send_request(
|
||||
ClientRequest(
|
||||
CallToolRequest(
|
||||
params=CallToolRequestParams(name="test_tool", arguments={}),
|
||||
)
|
||||
),
|
||||
CallToolResult,
|
||||
)
|
||||
|
||||
# Second request (should work normally)
|
||||
result = await client.send_request(
|
||||
ClientRequest(
|
||||
CallToolRequest(
|
||||
params=CallToolRequestParams(name="test_tool", arguments={}),
|
||||
)
|
||||
),
|
||||
CallToolResult,
|
||||
)
|
||||
|
||||
# Verify second request completed successfully
|
||||
assert len(result.content) == 1
|
||||
# Type narrowing for pyright
|
||||
content = result.content[0]
|
||||
assert content.type == "text"
|
||||
assert isinstance(content, types.TextContent)
|
||||
assert content.text == "Call number: 2"
|
||||
assert call_count == 2
|
||||
# Verify second request completed successfully
|
||||
assert len(result.content) == 1
|
||||
# Type narrowing for pyright
|
||||
content = result.content[0]
|
||||
assert content.type == "text"
|
||||
assert isinstance(content, types.TextContent)
|
||||
assert content.text == "Call number: 2"
|
||||
assert call_count == 2
|
||||
|
|
|
|||
|
|
@ -6,8 +6,8 @@ from typing import Any
|
|||
|
||||
import pytest
|
||||
|
||||
from mcp import Client
|
||||
from mcp.server.lowlevel import Server
|
||||
from mcp.shared.memory import create_connected_server_and_client_session
|
||||
from mcp.types import (
|
||||
Completion,
|
||||
CompletionArgument,
|
||||
|
|
@ -38,7 +38,7 @@ async def test_completion_handler_receives_context():
|
|||
# Return test completion
|
||||
return Completion(values=["test-completion"], total=1, has_more=False)
|
||||
|
||||
async with create_connected_server_and_client_session(server) as client:
|
||||
async with Client(server) as client:
|
||||
# Test with context
|
||||
result = await client.complete(
|
||||
ref=ResourceTemplateReference(type="ref/resource", uri="test://resource/{param}"),
|
||||
|
|
@ -70,7 +70,7 @@ async def test_completion_backward_compatibility():
|
|||
|
||||
return Completion(values=["no-context-completion"], total=1, has_more=False)
|
||||
|
||||
async with create_connected_server_and_client_session(server) as client:
|
||||
async with Client(server) as client:
|
||||
# Test without context
|
||||
result = await client.complete(
|
||||
ref=PromptReference(type="ref/prompt", name="test-prompt"), argument={"name": "arg", "value": "val"}
|
||||
|
|
@ -109,7 +109,7 @@ async def test_dependent_completion_scenario():
|
|||
|
||||
return Completion(values=[], total=0, has_more=False) # pragma: no cover
|
||||
|
||||
async with create_connected_server_and_client_session(server) as client:
|
||||
async with Client(server) as client:
|
||||
# First, complete database
|
||||
db_result = await client.complete(
|
||||
ref=ResourceTemplateReference(type="ref/resource", uri="db://{database}/{table}"),
|
||||
|
|
@ -160,7 +160,7 @@ async def test_completion_error_on_missing_context():
|
|||
|
||||
return Completion(values=[], total=0, has_more=False) # pragma: no cover
|
||||
|
||||
async with create_connected_server_and_client_session(server) as client:
|
||||
async with Client(server) as client:
|
||||
# Try to complete table without database context - should raise error
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
await client.complete(
|
||||
|
|
|
|||
Loading…
Reference in a new issue