import asyncio
import logging
import string
import pytest
from fastmcp import Client, FastMCP
from pymcp.middleware import ResponseMetadataMiddleware, StripUnknownArgumentsMiddleware
from pymcp.server import PyMCP
logger = logging.getLogger(__name__)
class TestStripUnknownArgumentsMiddleware:
"""Dedicated test class for the StripUnknownArgumentsMiddleware."""
@pytest.fixture(scope="class")
@classmethod
def mcp_server(cls):
"""Fixture to create an MCP server instance with the middleware."""
server = FastMCP()
mcp_obj = PyMCP()
server_with_features = mcp_obj.register_features(server)
server_with_features.add_middleware(StripUnknownArgumentsMiddleware())
return server_with_features
@pytest.fixture(scope="class", autouse=True)
@classmethod
def mcp_client(cls, mcp_server):
"""Fixture to create a client for the MCP server."""
mcp_client = Client(transport=mcp_server, timeout=60)
return mcp_client
async def call_tool(self, tool_name: str, mcp_client: Client, **kwargs):
"""Helper method to call a tool on the MCP server."""
async with mcp_client:
result = await mcp_client.call_tool(tool_name, arguments=kwargs)
await mcp_client.close()
return result
def test_strip_unknown_arguments(self, mcp_client: Client, caplog):
"""Test that unknown arguments are stripped from tool calls and logged."""
tool_name = "greet"
valid_name = "Test User"
unknown_arg_value = "This should be stripped"
with caplog.at_level(logging.INFO):
results = asyncio.run(
self.call_tool(
tool_name,
mcp_client,
name=valid_name,
unknown_argument=unknown_arg_value,
)
)
# Verify the tool call succeeded with valid argument
assert hasattr(results, "content"), "Expected results to have 'content' attribute"
assert hasattr(results, "structured_content"), "Expected results to have 'structured_content' attribute"
assert "result" in results.structured_content, "Expected 'structured_content' to have 'result' key"
# Verify the greeting contains the valid name (proving valid args passed through)
greeting = results.structured_content["result"]
assert valid_name in greeting, f"Expected greeting to contain '{valid_name}'"
# Verify logging occurred for unknown arguments
assert any("Unknown arguments for tool 'greet'" in record.message for record in caplog.records), (
"Expected logging of unknown arguments"
)
# Verify the unknown argument was identified in the logs
assert any("unknown_argument" in record.message for record in caplog.records), (
"Expected 'unknown_argument' to be logged as unknown"
)
def test_all_arguments_unknown(self, mcp_client: Client, caplog):
"""Test behavior when all provided arguments are unknown."""
tool_name = "greet"
with caplog.at_level(logging.INFO):
results = asyncio.run(
self.call_tool(
tool_name,
mcp_client,
completely_unknown_arg="value1",
another_unknown_arg="value2",
)
)
# Verify the tool call still succeeds (using defaults)
assert hasattr(results, "content"), "Expected results to have 'content' attribute"
assert hasattr(results, "structured_content"), "Expected results to have 'structured_content' attribute"
assert "result" in results.structured_content, "Expected 'structured_content' to have 'result' key"
# Verify default greeting (no name provided)
greeting = results.structured_content["result"]
assert "World" in greeting, "Expected default greeting with 'World'"
# Verify logging occurred
assert any("Unknown arguments for tool 'greet'" in record.message for record in caplog.records), (
"Expected logging of unknown arguments"
)
def test_no_arguments_provided(self, mcp_client: Client, caplog):
"""Test that middleware handles tools called with no arguments correctly."""
tool_name = "greet"
with caplog.at_level(logging.INFO):
results = asyncio.run(self.call_tool(tool_name, mcp_client))
# Verify the tool call succeeds
assert hasattr(results, "content"), "Expected results to have 'content' attribute"
assert hasattr(results, "structured_content"), "Expected results to have 'structured_content' attribute"
# Verify no middleware logging for this case (no args to strip)
middleware_logs = [record for record in caplog.records if "Unknown arguments" in record.message]
assert len(middleware_logs) == 0, "Expected no middleware logging when no arguments provided"
def test_only_valid_arguments(self, mcp_client: Client, caplog):
"""Test that middleware doesn't interfere when only valid arguments are provided."""
tool_name = "greet"
valid_name = "Valid User"
with caplog.at_level(logging.INFO):
results = asyncio.run(self.call_tool(tool_name, mcp_client, name=valid_name))
# Verify the tool call succeeds with the valid argument
assert hasattr(results, "content"), "Expected results to have 'content' attribute"
greeting = results.structured_content["result"]
assert valid_name in greeting, f"Expected greeting to contain '{valid_name}'"
# Verify no unknown argument logging
unknown_arg_logs = [record for record in caplog.records if "Unknown arguments" in record.message]
assert len(unknown_arg_logs) == 0, "Expected no unknown argument logging for valid args only"
def test_mixed_valid_and_unknown_arguments(self, mcp_client: Client, caplog):
"""Test middleware behavior with a mix of valid and unknown arguments."""
tool_name = "greet"
valid_name = "Mixed Test"
with caplog.at_level(logging.INFO):
results = asyncio.run(
self.call_tool(
tool_name,
mcp_client,
name=valid_name,
unknown1="value1",
unknown2={"key": "value2"},
unknown3=3.14,
)
)
# Verify valid argument was used
greeting = results.structured_content["result"]
assert valid_name in greeting, f"Expected greeting to contain '{valid_name}'"
# Verify multiple unknown arguments are logged
unknown_logs = [record for record in caplog.records if "Unknown arguments for tool 'greet'" in record.message]
assert len(unknown_logs) > 0, "Expected logging for unknown arguments"
# Verify all three unknown arguments are mentioned
log_messages = " ".join([record.message for record in caplog.records])
assert "unknown1" in log_messages, "Expected 'unknown1' in logs"
assert "unknown2" in log_messages, "Expected 'unknown2' in logs"
assert "unknown3" in log_messages, "Expected 'unknown3' in logs"
class TestResponseMetadataMiddleware:
"""Dedicated test class for the ResponseMetadataMiddleware."""
@pytest.fixture(scope="class")
@classmethod
def mcp_server(cls):
"""Fixture to create an MCP server instance with the middleware."""
server = FastMCP()
mcp_obj = PyMCP()
server_with_features = mcp_obj.register_features(server)
server_with_features.add_middleware(ResponseMetadataMiddleware())
return server_with_features
@pytest.fixture(scope="class", autouse=True)
@classmethod
def mcp_client(cls, mcp_server):
"""Fixture to create a client for the MCP server."""
mcp_client = Client(transport=mcp_server, timeout=60)
return mcp_client
async def call_tool(self, tool_name: str, mcp_client: Client, **kwargs):
"""Helper method to call a tool on the MCP server."""
async with mcp_client:
result = await mcp_client.call_tool(tool_name, arguments=kwargs)
await mcp_client.close()
return result
def test_call_for_package_metadata(self, mcp_client: Client, caplog):
"""Test that metadata is added to tool responses and appropriate logging occurs."""
tool_name = "greet"
valid_name = "Test User"
with caplog.at_level(logging.DEBUG):
results = asyncio.run(
self.call_tool(
tool_name,
mcp_client,
name=valid_name,
)
)
# Verify the tool call succeeded with valid argument
assert hasattr(results, "content"), "Expected results to have 'content' attribute"
assert hasattr(results, "structured_content"), "Expected results to have 'structured_content' attribute"
assert "result" in results.structured_content, "Expected 'structured_content' to have 'result' key"
# Verify the greeting contains the valid name (proving valid args passed through)
greeting = results.structured_content["result"]
assert valid_name in greeting, f"Expected greeting to contain '{valid_name}'"
assert getattr(results, "meta", None) is not None, "Expected results to have a valid 'meta' attribute"
assert ResponseMetadataMiddleware.PACKAGE_METADATA_KEY in results.meta, (
f"Expected '{ResponseMetadataMiddleware.PACKAGE_METADATA_KEY}' in meta"
)
assert ResponseMetadataMiddleware.TIMING_METADATA_KEY in results.meta, (
f"Expected '{ResponseMetadataMiddleware.TIMING_METADATA_KEY}' in meta"
)
assert "name" in results.meta[ResponseMetadataMiddleware.PACKAGE_METADATA_KEY], (
"Expected 'name' in package metadata"
)
assert "version" in results.meta[ResponseMetadataMiddleware.PACKAGE_METADATA_KEY], (
"Expected 'version' in package metadata"
)
assert results.meta[ResponseMetadataMiddleware.PACKAGE_METADATA_KEY]["name"] == "pymcp-template"
assert "tool_response_time_ms" in results.meta[ResponseMetadataMiddleware.TIMING_METADATA_KEY], (
"Expected 'tool_response_time_ms' in timing metadata"
)
assert isinstance(
results.meta[ResponseMetadataMiddleware.TIMING_METADATA_KEY]["tool_response_time_ms"], float
), "Expected 'tool_response_time_ms' to be a float"
# Verify logging occurred for metadata addition
assert any("Added package metadata to tool response" in record.message for record in caplog.records), (
"Expected debug logging of package metadata addition"
)
def test_call_tool_no_response(self, mcp_client: Client, caplog):
"""Test that the middleware passes up exceptions."""
tool_name = "permutations"
with caplog.at_level(logging.DEBUG):
try:
results = asyncio.run(
self.call_tool(
tool_name,
mcp_client,
n=5,
k=6, # k > n to force exception
)
)
except Exception as e:
logger.error(f"Exception during tool call. {e}", exc_info=True)
results = None
# Verify the tool call returns None because an exception was raised (k > n is invalid)
assert results is None, "Expected results to be None because of exception in tool"
# Verify error logging occurred
assert any(
record.levelno == logging.ERROR and "cannot be greater" in record.message for record in caplog.records
), "Expected error logging due to exception in tool call"
assert any("failed after" in record.message for record in caplog.records), (
"Expected warning logging of operation failure"
)
# Verify no logging occurred for metadata addition since result is None
assert not any("Added package metadata to tool response" in record.message for record in caplog.records), (
"Did not expect debug logging of package metadata addition when exceptions occur"
)
def test_call_tool_with_specific_metadata(self, mcp_client: Client, caplog):
"""Test that existing metadata is preserved and package metadata is added."""
tool_name = "generate_password"
expected_password_length = 12
with caplog.at_level(logging.DEBUG):
results = asyncio.run(
self.call_tool(
tool_name,
mcp_client,
use_special_chars=True,
length=expected_password_length,
)
)
# Verify the tool call succeeded with valid argument
assert hasattr(results, "content"), "Expected results to have 'content' attribute"
assert hasattr(results, "structured_content"), "Expected results to have 'structured_content' attribute"
assert "result" in results.structured_content, "Expected 'structured_content' to have 'result' key"
# Verify the generated password has the expected length
generated_password = results.structured_content["result"]
assert len(generated_password) == expected_password_length, (
f"Expected generated password to be of length {expected_password_length}, got {len(generated_password)}"
)
assert getattr(results, "meta", None) is not None, "Expected results to have a valid 'meta' attribute"
assert ResponseMetadataMiddleware.PACKAGE_METADATA_KEY in results.meta, (
f"Expected '{ResponseMetadataMiddleware.PACKAGE_METADATA_KEY}' in meta"
)
assert ResponseMetadataMiddleware.TIMING_METADATA_KEY in results.meta, (
f"Expected '{ResponseMetadataMiddleware.TIMING_METADATA_KEY}' in meta"
)
assert "name" in results.meta[ResponseMetadataMiddleware.PACKAGE_METADATA_KEY], (
"Expected 'name' in package metadata"
)
assert "version" in results.meta[ResponseMetadataMiddleware.PACKAGE_METADATA_KEY], (
"Expected 'version' in package metadata"
)
assert results.meta[ResponseMetadataMiddleware.PACKAGE_METADATA_KEY]["name"] == "pymcp-template"
assert "tool_response_time_ms" in results.meta[ResponseMetadataMiddleware.TIMING_METADATA_KEY], (
"Expected 'tool_response_time_ms' in timing metadata"
)
assert isinstance(
results.meta[ResponseMetadataMiddleware.TIMING_METADATA_KEY]["tool_response_time_ms"], float
), "Expected 'tool_response_time_ms' to be a float"
# Verify tool specific metadata is still present and values make sense
assert "length_satisfied" in results.meta[tool_name], (
"Expected tool specific metadata key 'length_satisfied' to be present"
)
assert results.meta[tool_name]["length_satisfied"] is True, (
"Expected 'length_satisfied' to be True since password matches requested length"
)
assert "generation_attempts" in results.meta[tool_name], (
"Expected tool specific metadata key 'generation_attempts' to be present"
)
assert results.meta[tool_name]["generation_attempts"] >= 1, "Expected 'generation_attempts' to be at least 1"
assert "character_set" in results.meta[tool_name], (
"Expected tool specific metadata key 'character_set' to be present"
)
assert isinstance(results.meta[tool_name]["character_set"], str), "Expected 'character_set' to be a string"
# Verify character_set contains expected character types when use_special_chars=True
# The character_set should be the pool of available characters for password generation
character_set = results.meta[tool_name]["character_set"]
assert any(c in character_set for c in string.ascii_lowercase), (
"Expected character_set to contain lowercase letters"
)
assert any(c in character_set for c in string.ascii_uppercase), (
"Expected character_set to contain uppercase letters"
)
assert any(c in character_set for c in string.digits), "Expected character_set to contain digits"
assert any(c in character_set for c in string.punctuation), (
"Expected character_set to contain punctuation when use_special_chars=True"
)
# Verify logging occurred for metadata addition
assert any("Added package metadata to tool response" in record.message for record in caplog.records), (
"Expected debug logging of package metadata addition"
)
