mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-09 15:19:10 +02:00
602 lines
22 KiB
Python
602 lines
22 KiB
Python
"""Tests for error handling middleware."""
|
|
|
|
import logging
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
from mcp import McpError
|
|
|
|
from fastmcp.server.middleware.error_handling import (
|
|
ErrorHandlingMiddleware,
|
|
RetryMiddleware,
|
|
)
|
|
from fastmcp.server.middleware.middleware import MiddlewareContext
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_context():
|
|
"""Create a mock middleware context."""
|
|
context = MagicMock(spec=MiddlewareContext)
|
|
context.method = "test_method"
|
|
return context
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_call_next():
|
|
"""Create a mock call_next function."""
|
|
return AsyncMock(return_value="test_result")
|
|
|
|
|
|
class TestErrorHandlingMiddleware:
|
|
"""Test error handling middleware functionality."""
|
|
|
|
def test_init_default(self):
|
|
"""Test default initialization."""
|
|
middleware = ErrorHandlingMiddleware()
|
|
assert middleware.logger.name == "fastmcp.errors"
|
|
assert middleware.include_traceback is False
|
|
assert middleware.error_callback is None
|
|
assert middleware.transform_errors is True
|
|
assert middleware.error_counts == {}
|
|
|
|
def test_init_custom(self):
|
|
"""Test custom initialization."""
|
|
logger = logging.getLogger("custom")
|
|
callback = MagicMock()
|
|
|
|
middleware = ErrorHandlingMiddleware(
|
|
logger=logger,
|
|
include_traceback=True,
|
|
error_callback=callback,
|
|
transform_errors=False,
|
|
)
|
|
assert middleware.logger is logger
|
|
assert middleware.include_traceback is True
|
|
assert middleware.error_callback is callback
|
|
assert middleware.transform_errors is False
|
|
|
|
def test_log_error_basic(self, mock_context, caplog):
|
|
"""Test basic error logging."""
|
|
middleware = ErrorHandlingMiddleware()
|
|
error = ValueError("test error")
|
|
|
|
with caplog.at_level(logging.ERROR):
|
|
middleware._log_error(error, mock_context)
|
|
|
|
assert "Error in test_method: ValueError: test error" in caplog.text
|
|
assert "ValueError:test_method" in middleware.error_counts
|
|
assert middleware.error_counts["ValueError:test_method"] == 1
|
|
|
|
def test_log_error_with_traceback(self, mock_context, caplog):
|
|
"""Test error logging with traceback."""
|
|
middleware = ErrorHandlingMiddleware(include_traceback=True)
|
|
error = ValueError("test error")
|
|
|
|
with caplog.at_level(logging.ERROR):
|
|
middleware._log_error(error, mock_context)
|
|
|
|
assert "Error in test_method: ValueError: test error" in caplog.text
|
|
# The traceback is added to the log message
|
|
assert "Error in test_method: ValueError: test error" in caplog.text
|
|
|
|
def test_log_error_with_callback(self, mock_context):
|
|
"""Test error logging with callback."""
|
|
callback = MagicMock()
|
|
middleware = ErrorHandlingMiddleware(error_callback=callback)
|
|
error = ValueError("test error")
|
|
|
|
middleware._log_error(error, mock_context)
|
|
|
|
callback.assert_called_once_with(error, mock_context)
|
|
|
|
def test_log_error_callback_exception(self, mock_context, caplog):
|
|
"""Test error logging when callback raises exception."""
|
|
callback = MagicMock(side_effect=RuntimeError("callback error"))
|
|
middleware = ErrorHandlingMiddleware(error_callback=callback)
|
|
error = ValueError("test error")
|
|
|
|
with caplog.at_level(logging.ERROR):
|
|
middleware._log_error(error, mock_context)
|
|
|
|
assert "Error in error callback: callback error" in caplog.text
|
|
|
|
def test_transform_error_mcp_error(self):
|
|
"""Test that MCP errors are not transformed."""
|
|
middleware = ErrorHandlingMiddleware()
|
|
from mcp.types import ErrorData
|
|
|
|
error = McpError(ErrorData(code=-32001, message="test error"))
|
|
|
|
result = middleware._transform_error(error)
|
|
|
|
assert result is error
|
|
|
|
def test_transform_error_disabled(self):
|
|
"""Test error transformation when disabled."""
|
|
middleware = ErrorHandlingMiddleware(transform_errors=False)
|
|
error = ValueError("test error")
|
|
|
|
result = middleware._transform_error(error)
|
|
|
|
assert result is error
|
|
|
|
def test_transform_error_value_error(self):
|
|
"""Test transforming ValueError."""
|
|
middleware = ErrorHandlingMiddleware()
|
|
error = ValueError("test error")
|
|
|
|
result = middleware._transform_error(error)
|
|
|
|
assert isinstance(result, McpError)
|
|
assert result.error.code == -32602
|
|
assert "Invalid params: test error" in result.error.message
|
|
|
|
def test_transform_error_file_not_found(self):
|
|
"""Test transforming FileNotFoundError."""
|
|
middleware = ErrorHandlingMiddleware()
|
|
error = FileNotFoundError("test error")
|
|
|
|
result = middleware._transform_error(error)
|
|
|
|
assert isinstance(result, McpError)
|
|
assert result.error.code == -32001
|
|
assert "Resource not found: test error" in result.error.message
|
|
|
|
def test_transform_error_permission_error(self):
|
|
"""Test transforming PermissionError."""
|
|
middleware = ErrorHandlingMiddleware()
|
|
error = PermissionError("test error")
|
|
|
|
result = middleware._transform_error(error)
|
|
|
|
assert isinstance(result, McpError)
|
|
assert result.error.code == -32000
|
|
assert "Permission denied: test error" in result.error.message
|
|
|
|
def test_transform_error_timeout_error(self):
|
|
"""Test transforming TimeoutError."""
|
|
middleware = ErrorHandlingMiddleware()
|
|
error = TimeoutError("test error")
|
|
|
|
result = middleware._transform_error(error)
|
|
|
|
assert isinstance(result, McpError)
|
|
assert result.error.code == -32000
|
|
assert "Request timeout: test error" in result.error.message
|
|
|
|
def test_transform_error_generic(self):
|
|
"""Test transforming generic error."""
|
|
middleware = ErrorHandlingMiddleware()
|
|
error = RuntimeError("test error")
|
|
|
|
result = middleware._transform_error(error)
|
|
|
|
assert isinstance(result, McpError)
|
|
assert result.error.code == -32603
|
|
assert "Internal error: test error" in result.error.message
|
|
|
|
async def test_on_message_success(self, mock_context, mock_call_next):
|
|
"""Test successful message handling."""
|
|
middleware = ErrorHandlingMiddleware()
|
|
|
|
result = await middleware.on_message(mock_context, mock_call_next)
|
|
|
|
assert result == "test_result"
|
|
assert mock_call_next.called
|
|
|
|
async def test_on_message_error_transform(self, mock_context, caplog):
|
|
"""Test error handling with transformation."""
|
|
middleware = ErrorHandlingMiddleware()
|
|
mock_call_next = AsyncMock(side_effect=ValueError("test error"))
|
|
|
|
with caplog.at_level(logging.ERROR):
|
|
with pytest.raises(McpError) as exc_info:
|
|
await middleware.on_message(mock_context, mock_call_next)
|
|
|
|
assert isinstance(exc_info.value, McpError)
|
|
assert exc_info.value.error.code == -32602
|
|
assert "Invalid params: test error" in exc_info.value.error.message
|
|
assert "Error in test_method: ValueError: test error" in caplog.text
|
|
|
|
def test_get_error_stats(self, mock_context):
|
|
"""Test getting error statistics."""
|
|
middleware = ErrorHandlingMiddleware()
|
|
error1 = ValueError("error1")
|
|
error2 = ValueError("error2")
|
|
error3 = RuntimeError("error3")
|
|
|
|
middleware._log_error(error1, mock_context)
|
|
middleware._log_error(error2, mock_context)
|
|
middleware._log_error(error3, mock_context)
|
|
|
|
stats = middleware.get_error_stats()
|
|
assert stats["ValueError:test_method"] == 2
|
|
assert stats["RuntimeError:test_method"] == 1
|
|
|
|
|
|
class TestRetryMiddleware:
|
|
"""Test retry middleware functionality."""
|
|
|
|
def test_init_default(self):
|
|
"""Test default initialization."""
|
|
middleware = RetryMiddleware()
|
|
assert middleware.max_retries == 3
|
|
assert middleware.base_delay == 1.0
|
|
assert middleware.max_delay == 60.0
|
|
assert middleware.backoff_multiplier == 2.0
|
|
assert middleware.retry_exceptions == (ConnectionError, TimeoutError)
|
|
assert middleware.logger.name == "fastmcp.retry"
|
|
|
|
def test_init_custom(self):
|
|
"""Test custom initialization."""
|
|
logger = logging.getLogger("custom")
|
|
middleware = RetryMiddleware(
|
|
max_retries=5,
|
|
base_delay=2.0,
|
|
max_delay=120.0,
|
|
backoff_multiplier=3.0,
|
|
retry_exceptions=(ValueError, RuntimeError),
|
|
logger=logger,
|
|
)
|
|
assert middleware.max_retries == 5
|
|
assert middleware.base_delay == 2.0
|
|
assert middleware.max_delay == 120.0
|
|
assert middleware.backoff_multiplier == 3.0
|
|
assert middleware.retry_exceptions == (ValueError, RuntimeError)
|
|
assert middleware.logger is logger
|
|
|
|
def test_should_retry_true(self):
|
|
"""Test retry decision for retryable errors."""
|
|
middleware = RetryMiddleware()
|
|
|
|
assert middleware._should_retry(ConnectionError()) is True
|
|
assert middleware._should_retry(TimeoutError()) is True
|
|
|
|
def test_should_retry_false(self):
|
|
"""Test retry decision for non-retryable errors."""
|
|
middleware = RetryMiddleware()
|
|
|
|
assert middleware._should_retry(ValueError()) is False
|
|
assert middleware._should_retry(RuntimeError()) is False
|
|
|
|
def test_calculate_delay(self):
|
|
"""Test delay calculation."""
|
|
middleware = RetryMiddleware(
|
|
base_delay=1.0, backoff_multiplier=2.0, max_delay=10.0
|
|
)
|
|
|
|
assert middleware._calculate_delay(0) == 1.0
|
|
assert middleware._calculate_delay(1) == 2.0
|
|
assert middleware._calculate_delay(2) == 4.0
|
|
assert middleware._calculate_delay(3) == 8.0
|
|
assert middleware._calculate_delay(4) == 10.0 # capped at max_delay
|
|
|
|
async def test_on_request_success_first_try(self, mock_context, mock_call_next):
|
|
"""Test successful request on first try."""
|
|
middleware = RetryMiddleware()
|
|
|
|
result = await middleware.on_request(mock_context, mock_call_next)
|
|
|
|
assert result == "test_result"
|
|
assert mock_call_next.call_count == 1
|
|
|
|
async def test_on_request_success_after_retries(self, mock_context, caplog):
|
|
"""Test successful request after retries."""
|
|
middleware = RetryMiddleware(base_delay=0.01) # Fast retry for testing
|
|
|
|
# Fail first two attempts, succeed on third
|
|
mock_call_next = AsyncMock(
|
|
side_effect=[
|
|
ConnectionError("connection failed"),
|
|
ConnectionError("connection failed"),
|
|
"test_result",
|
|
]
|
|
)
|
|
|
|
with caplog.at_level(logging.WARNING):
|
|
result = await middleware.on_request(mock_context, mock_call_next)
|
|
|
|
assert result == "test_result"
|
|
assert mock_call_next.call_count == 3
|
|
assert "Retrying in" in caplog.text
|
|
|
|
async def test_on_request_max_retries_exceeded(self, mock_context, caplog):
|
|
"""Test request failing after max retries."""
|
|
middleware = RetryMiddleware(max_retries=2, base_delay=0.01)
|
|
|
|
# Fail all attempts
|
|
mock_call_next = AsyncMock(side_effect=ConnectionError("connection failed"))
|
|
|
|
with caplog.at_level(logging.WARNING):
|
|
with pytest.raises(ConnectionError):
|
|
await middleware.on_request(mock_context, mock_call_next)
|
|
|
|
assert mock_call_next.call_count == 3 # initial + 2 retries
|
|
assert "Retrying in" in caplog.text
|
|
|
|
async def test_on_request_non_retryable_error(self, mock_context):
|
|
"""Test non-retryable error is not retried."""
|
|
middleware = RetryMiddleware()
|
|
mock_call_next = AsyncMock(side_effect=ValueError("non-retryable"))
|
|
|
|
with pytest.raises(ValueError):
|
|
await middleware.on_request(mock_context, mock_call_next)
|
|
|
|
assert mock_call_next.call_count == 1 # No retries
|
|
|
|
|
|
@pytest.fixture
|
|
def error_handling_server():
|
|
"""Create a FastMCP server specifically for error handling middleware tests."""
|
|
from fastmcp import FastMCP
|
|
|
|
mcp = FastMCP("ErrorHandlingTestServer")
|
|
|
|
@mcp.tool
|
|
def reliable_operation(data: str) -> str:
|
|
"""A reliable operation that always succeeds."""
|
|
return f"Success: {data}"
|
|
|
|
@mcp.tool
|
|
def failing_operation(error_type: str = "value") -> str:
|
|
"""An operation that fails with different error types."""
|
|
if error_type == "value":
|
|
raise ValueError("Value error occurred")
|
|
elif error_type == "file":
|
|
raise FileNotFoundError("File not found")
|
|
elif error_type == "permission":
|
|
raise PermissionError("Permission denied")
|
|
elif error_type == "timeout":
|
|
raise TimeoutError("Operation timed out")
|
|
elif error_type == "generic":
|
|
raise RuntimeError("Generic runtime error")
|
|
else:
|
|
return "Operation completed"
|
|
|
|
@mcp.tool
|
|
def intermittent_operation(fail_rate: float = 0.5) -> str:
|
|
"""An operation that fails intermittently."""
|
|
import random
|
|
|
|
if random.random() < fail_rate:
|
|
raise ConnectionError("Random connection failure")
|
|
return "Operation succeeded"
|
|
|
|
@mcp.tool
|
|
def retryable_operation(attempt_count: int = 0) -> str:
|
|
"""An operation that succeeds after a few attempts."""
|
|
# This is a simple way to simulate retry behavior
|
|
# In a real scenario, you might use external state
|
|
if attempt_count < 2:
|
|
raise ConnectionError("Temporary connection error")
|
|
return "Operation succeeded after retries"
|
|
|
|
return mcp
|
|
|
|
|
|
class TestErrorHandlingMiddlewareIntegration:
|
|
"""Integration tests for error handling middleware with real FastMCP server."""
|
|
|
|
async def test_error_handling_middleware_logs_real_errors(
|
|
self, error_handling_server, caplog
|
|
):
|
|
"""Test that error handling middleware logs real errors from tools."""
|
|
from fastmcp.client import Client
|
|
|
|
error_handling_server.add_middleware(ErrorHandlingMiddleware())
|
|
|
|
with caplog.at_level(logging.ERROR):
|
|
async with Client(error_handling_server) as client:
|
|
# Test different types of errors
|
|
with pytest.raises(Exception):
|
|
await client.call_tool("failing_operation", {"error_type": "value"})
|
|
|
|
with pytest.raises(Exception):
|
|
await client.call_tool("failing_operation", {"error_type": "file"})
|
|
|
|
log_text = caplog.text
|
|
|
|
# Should have error logs for both failures
|
|
assert "Error in tools/call: ToolError:" in log_text
|
|
# Should have captured both error instances
|
|
error_count = log_text.count("Error in tools/call:")
|
|
assert error_count == 2
|
|
|
|
async def test_error_handling_middleware_tracks_error_statistics(
|
|
self, error_handling_server
|
|
):
|
|
"""Test that error handling middleware accurately tracks error statistics."""
|
|
from fastmcp.client import Client
|
|
|
|
error_middleware = ErrorHandlingMiddleware()
|
|
error_handling_server.add_middleware(error_middleware)
|
|
|
|
async with Client(error_handling_server) as client:
|
|
# Generate different types of errors
|
|
for _ in range(3):
|
|
with pytest.raises(Exception):
|
|
await client.call_tool("failing_operation", {"error_type": "value"})
|
|
|
|
for _ in range(2):
|
|
with pytest.raises(Exception):
|
|
await client.call_tool("failing_operation", {"error_type": "file"})
|
|
|
|
# Try some intermittent operations (some may succeed)
|
|
for _ in range(5):
|
|
try:
|
|
await client.call_tool("intermittent_operation", {"fail_rate": 0.8})
|
|
except Exception:
|
|
pass # Expected failures
|
|
|
|
# Check error statistics
|
|
stats = error_middleware.get_error_stats()
|
|
|
|
# Should have tracked the ToolError wrapper
|
|
assert "ToolError:tools/call" in stats
|
|
assert stats["ToolError:tools/call"] >= 5 # At least the 5 deliberate failures
|
|
|
|
async def test_error_handling_middleware_with_success_and_failure(
|
|
self, error_handling_server, caplog
|
|
):
|
|
"""Test error handling middleware with mix of successful and failed operations."""
|
|
from fastmcp.client import Client
|
|
|
|
error_handling_server.add_middleware(ErrorHandlingMiddleware())
|
|
|
|
with caplog.at_level(logging.ERROR):
|
|
async with Client(error_handling_server) as client:
|
|
# Successful operation (should not generate error logs)
|
|
await client.call_tool("reliable_operation", {"data": "test"})
|
|
|
|
# Failed operation (should generate error log)
|
|
with pytest.raises(Exception):
|
|
await client.call_tool("failing_operation", {"error_type": "value"})
|
|
|
|
# Another successful operation
|
|
await client.call_tool("reliable_operation", {"data": "test2"})
|
|
|
|
log_text = caplog.text
|
|
|
|
# Should only have one error log (for the failed operation)
|
|
error_count = log_text.count("Error in tools/call:")
|
|
assert error_count == 1
|
|
|
|
async def test_error_handling_middleware_custom_callback(
|
|
self, error_handling_server
|
|
):
|
|
"""Test error handling middleware with custom error callback."""
|
|
from fastmcp.client import Client
|
|
|
|
captured_errors = []
|
|
|
|
def error_callback(error, context):
|
|
captured_errors.append(
|
|
{
|
|
"error_type": type(error).__name__,
|
|
"message": str(error),
|
|
"method": context.method,
|
|
}
|
|
)
|
|
|
|
error_handling_server.add_middleware(
|
|
ErrorHandlingMiddleware(error_callback=error_callback)
|
|
)
|
|
|
|
async with Client(error_handling_server) as client:
|
|
# Generate some errors
|
|
with pytest.raises(Exception):
|
|
await client.call_tool("failing_operation", {"error_type": "value"})
|
|
|
|
with pytest.raises(Exception):
|
|
await client.call_tool("failing_operation", {"error_type": "timeout"})
|
|
|
|
# Check that callback was called
|
|
assert len(captured_errors) == 2
|
|
assert captured_errors[0]["error_type"] == "ToolError"
|
|
assert captured_errors[1]["error_type"] == "ToolError"
|
|
assert all(error["method"] == "tools/call" for error in captured_errors)
|
|
|
|
async def test_error_handling_middleware_transform_errors(
|
|
self, error_handling_server
|
|
):
|
|
"""Test error transformation functionality."""
|
|
from fastmcp.client import Client
|
|
|
|
error_handling_server.add_middleware(
|
|
ErrorHandlingMiddleware(transform_errors=True)
|
|
)
|
|
|
|
async with Client(error_handling_server) as client:
|
|
# All errors should still be raised, but potentially transformed
|
|
with pytest.raises(Exception) as exc_info:
|
|
await client.call_tool("failing_operation", {"error_type": "value"})
|
|
|
|
# Error should still exist (may be wrapped by FastMCP)
|
|
assert exc_info.value is not None
|
|
|
|
|
|
class TestRetryMiddlewareIntegration:
|
|
"""Integration tests for retry middleware with real FastMCP server."""
|
|
|
|
async def test_retry_middleware_with_transient_failures(
|
|
self, error_handling_server, caplog
|
|
):
|
|
"""Test retry middleware with operations that have transient failures."""
|
|
from fastmcp.client import Client
|
|
|
|
# Configure retry middleware to retry connection errors
|
|
error_handling_server.add_middleware(
|
|
RetryMiddleware(
|
|
max_retries=3,
|
|
base_delay=0.01, # Very short delay for testing
|
|
retry_exceptions=(ConnectionError,),
|
|
)
|
|
)
|
|
|
|
with caplog.at_level(logging.WARNING):
|
|
async with Client(error_handling_server) as client:
|
|
# This operation fails intermittently - try several times
|
|
success_count = 0
|
|
for _ in range(5):
|
|
try:
|
|
await client.call_tool(
|
|
"intermittent_operation", {"fail_rate": 0.7}
|
|
)
|
|
success_count += 1
|
|
except Exception:
|
|
pass # Some failures expected even with retries
|
|
|
|
# Should have some retry log messages
|
|
# Note: Retry logs might not appear if the underlying errors are wrapped by FastMCP
|
|
# The key is that some operations should succeed due to retries
|
|
|
|
async def test_retry_middleware_with_permanent_failures(
|
|
self, error_handling_server
|
|
):
|
|
"""Test that retry middleware doesn't retry non-retryable errors."""
|
|
from fastmcp.client import Client
|
|
|
|
# Configure retry middleware for connection errors only
|
|
error_handling_server.add_middleware(
|
|
RetryMiddleware(
|
|
max_retries=3, base_delay=0.01, retry_exceptions=(ConnectionError,)
|
|
)
|
|
)
|
|
|
|
async with Client(error_handling_server) as client:
|
|
# Value errors should not be retried
|
|
with pytest.raises(Exception):
|
|
await client.call_tool("failing_operation", {"error_type": "value"})
|
|
|
|
# Should fail immediately without retries
|
|
|
|
async def test_combined_error_handling_and_retry_middleware(
|
|
self, error_handling_server, caplog
|
|
):
|
|
"""Test error handling and retry middleware working together."""
|
|
from fastmcp.client import Client
|
|
|
|
# Add both middleware
|
|
error_handling_server.add_middleware(ErrorHandlingMiddleware())
|
|
error_handling_server.add_middleware(
|
|
RetryMiddleware(
|
|
max_retries=2, base_delay=0.01, retry_exceptions=(ConnectionError,)
|
|
)
|
|
)
|
|
|
|
with caplog.at_level(logging.ERROR):
|
|
async with Client(error_handling_server) as client:
|
|
# Try intermittent operation
|
|
try:
|
|
await client.call_tool("intermittent_operation", {"fail_rate": 0.9})
|
|
except Exception:
|
|
pass # May still fail even with retries
|
|
|
|
# Try permanent failure
|
|
with pytest.raises(Exception):
|
|
await client.call_tool("failing_operation", {"error_type": "value"})
|
|
|
|
log_text = caplog.text
|
|
|
|
# Should have error logs from error handling middleware
|
|
assert "Error in tools/call:" in log_text
|