ceos/fastmcp
1
0
Fork 0
mirror of https://github.com/PrefectHQ/fastmcp.git synced 2026-08-24 14:34:17 +02:00
Code Issues Projects Releases Packages Wiki Activity

Update middleware logging snapshots and rate-limit tests for SDK v2 (param shapes, MCPError surface)

Browse source
This commit is contained in:
Jeremiah Lowin 2026-07-05 23:03:47 -04:00
commit b0adc8eb9c
No known key found for this signature in database
2 changed files with 70 additions and 49 deletions
Download patch file Download diff file Expand all files Collapse all files

21
tests/server/middleware/test_logging.py
View file

@ -9,7 +9,6 @@ from unittest.mock import AsyncMock, MagicMock, patch
import mcp_types
import pytest
from inline_snapshot import snapshot
from pydantic import AnyUrl
from fastmcp import FastMCP
from fastmcp.client import Client
@ -142,7 +141,7 @@ class TestStructuredLoggingMiddleware:
"event": "request_start",
"source": "client",
"method": "test_method",
"payload": '{"method":"tools/call","params":{"task":null,"_meta":null,"name":"test_method","arguments":{"param":"value"}}}',
"payload": '{"method":"tools/call","params":{"_meta":null,"inputResponses":null,"requestState":null,"name":"test_method","arguments":{"param":"value"},"task":null}}',
"payload_type": "CallToolRequest",
}
)
@ -157,7 +156,7 @@ class TestStructuredLoggingMiddleware:
"event": "request_start",
"source": "client",
"method": "test_method",
"payload_length": 110,
"payload_length": 152,
}
)
@ -175,8 +174,8 @@ class TestStructuredLoggingMiddleware:
"event": "request_start",
"source": "client",
"method": "test_method",
"payload_tokens": 27,
"payload_length": 110,
"payload_tokens": 38,
"payload_length": 152,
}
)
@ -284,7 +283,7 @@ class TestLoggingMiddleware:
message=mcp_types.ReadResourceRequest(
method="resources/read",
params=mcp_types.ReadResourceRequestParams(
uri=AnyUrl("test://example/1"),
uri="test://example/1",
),
)
)
@ -297,7 +296,7 @@ class TestLoggingMiddleware:
assert get_log_lines(caplog) == snapshot(
[
'{"event": "request_start", "method": "test_method", "source": "client", "payload": "{\\"method\\":\\"resources/read\\",\\"params\\":{\\"task\\":null,\\"_meta\\":null,\\"uri\\":\\"test://example/1\\"}}", "payload_type": "ReadResourceRequest"}',
'{"event": "request_start", "method": "test_method", "source": "client", "payload": "{\\"method\\":\\"resources/read\\",\\"params\\":{\\"_meta\\":null,\\"inputResponses\\":null,\\"requestState\\":null,\\"uri\\":\\"test://example/1\\"}}", "payload_type": "ReadResourceRequest"}',
'{"event": "request_success", "method": "test_method", "source": "client", "duration_ms": 0.02}',
]
)
@ -357,7 +356,7 @@ class TestLoggingMiddleware:
assert get_log_lines(caplog) == snapshot(
[
'{"event": "request_start", "method": "test_method", "source": "client", "payload": "{\\"method\\":\\"tools/call\\",\\"params\\":{\\"task\\":null,\\"_meta\\":null,\\"name\\":\\"test_method\\",\\"arguments\\":{\\"obj\\":\\"NON_SERIALIZABLE\\"}}}", "payload_type": "CallToolRequest"}',
'{"event": "request_start", "method": "test_method", "source": "client", "payload": "{\\"method\\":\\"tools/call\\",\\"params\\":{\\"_meta\\":null,\\"inputResponses\\":null,\\"requestState\\":null,\\"name\\":\\"test_method\\",\\"arguments\\":{\\"obj\\":\\"NON_SERIALIZABLE\\"},\\"task\\":null}}", "payload_type": "CallToolRequest"}',
'{"event": "request_success", "method": "test_method", "source": "client", "duration_ms": 0.02}',
]
)
@ -532,7 +531,7 @@ class TestLoggingMiddlewareIntegration:
assert get_log_lines(caplog) == snapshot(
[
'event=request_start method=tools/call source=client payload={"task":null,"_meta":null,"name":"simple_operation","arguments":{"data":"payload_test"}} payload_type=CallToolRequestParams',
'event=request_start method=tools/call source=client payload={"_meta":null,"inputResponses":null,"requestState":null,"name":"simple_operation","arguments":{"data":"payload_test"},"task":null} payload_type=CallToolRequestParams',
"event=request_success method=tools/call source=client duration_ms=0.02",
]
)
@ -555,7 +554,7 @@ class TestLoggingMiddlewareIntegration:
assert get_log_lines(caplog) == snapshot(
[
'{"event": "request_start", "method": "tools/call", "source": "client", "payload": "{\\"task\\":null,\\"_meta\\":null,\\"name\\":\\"simple_operation\\",\\"arguments\\":{\\"data\\":\\"json_test\\"}}", "payload_type": "CallToolRequestParams"}',
'{"event": "request_start", "method": "tools/call", "source": "client", "payload": "{\\"_meta\\":null,\\"inputResponses\\":null,\\"requestState\\":null,\\"name\\":\\"simple_operation\\",\\"arguments\\":{\\"data\\":\\"json_test\\"},\\"task\\":null}", "payload_type": "CallToolRequestParams"}',
'{"event": "request_success", "method": "tools/call", "source": "client", "duration_ms": 0.02}',
]
)
@ -648,6 +647,6 @@ class TestLoggingMiddlewareIntegration:
# Check that our custom logger captured the logs
log_output = log_buffer.getvalue()
assert log_output == snapshot("""\
event=request_start method=tools/call source=client payload={"task":null,"_meta":null,"name":"simple_operation","arguments":{"data":"custom_test"}} payload_type=CallToolRequestParams
event=request_start method=tools/call source=client payload={"_meta":null,"inputResponses":null,"requestState":null,"name":"simple_operation","arguments":{"data":"custom_test"},"task":null} payload_type=CallToolRequestParams
event=request_success method=tools/call source=client duration_ms=0.02
""")

98
tests/server/middleware/test_rate_limiting.py
View file

@ -4,10 +4,10 @@ import asyncio
from unittest.mock import AsyncMock, MagicMock
import pytest
from mcp import MCPError
from fastmcp import FastMCP
from fastmcp.client import Client
from fastmcp.exceptions import ToolError
from fastmcp.server.middleware.middleware import MiddlewareContext
from fastmcp.server.middleware.rate_limiting import (
RateLimitError,
@ -394,7 +394,7 @@ class TestRateLimitingMiddlewareIntegration:
for i in range(30):
try:
await client.call_tool("quick_action", {"message": str(i)})
except ToolError as exc:
except MCPError as exc:
assert "Rate limit exceeded" in str(exc)
hit_limit = True
break
@ -433,37 +433,53 @@ class TestRateLimitingMiddlewareIntegration:
async def test_sliding_window_rate_limiting(self, rate_limit_server):
"""Test sliding window rate limiting implementation."""
# A tight window; the SDK's internal init/list_tools requests count
# against it too, so fire tool calls until the limit is hit rather than
# hard-coding the exact request count (which differs across SDK versions).
rate_limit_server.add_middleware(
SlidingWindowRateLimitingMiddleware(
max_requests=6, # 1 init + 1 list_tools + 3 calls + 1 to fail
window_minutes=1, # 1-minute window
max_requests=6,
window_minutes=1,
)
)
async with Client(rate_limit_server) as client:
# Should allow up to the limit
await client.call_tool("quick_action", {"message": "1"})
await client.call_tool("quick_action", {"message": "2"})
await client.call_tool("quick_action", {"message": "3"})
# Fourth should be blocked
with pytest.raises(ToolError, match="Rate limit exceeded"):
await client.call_tool("quick_action", {"message": "4"})
hit_limit = False
for i in range(10):
try:
await client.call_tool("quick_action", {"message": str(i)})
except MCPError as exc:
assert "Rate limit exceeded" in str(exc)
hit_limit = True
break
assert hit_limit, "Rate limit was never triggered"
async def test_rate_limiting_with_different_operations(self, rate_limit_server):
"""Test that rate limiting applies to all types of operations."""
rate_limit_server.add_middleware(
RateLimitingMiddleware(max_requests_per_second=9.0, burst_capacity=5)
RateLimitingMiddleware(max_requests_per_second=0.001, burst_capacity=5)
)
async with Client(rate_limit_server) as client:
# Mix different operations
await client.call_tool("quick_action", {"message": "test"})
await client.call_tool("heavy_computation")
# Should be rate limited regardless of operation type
with pytest.raises(ToolError, match="Rate limit exceeded"):
await client.call_tool("batch_process", {"items": ["a", "b", "c"]})
# Mix different operations; with near-zero refill the shared bucket
# is exhausted regardless of which operation trips it.
operations = [
("quick_action", {"message": "test"}),
("heavy_computation", {}),
("batch_process", {"items": ["a", "b", "c"]}),
("quick_action", {"message": "again"}),
("heavy_computation", {}),
("batch_process", {"items": ["x"]}),
]
hit_limit = False
for name, args in operations:
try:
await client.call_tool(name, args)
except MCPError as exc:
assert "Rate limit exceeded" in str(exc)
hit_limit = True
break
assert hit_limit, "Rate limit was never triggered"
async def test_custom_client_identification(self, rate_limit_server):
"""Test rate limiting with custom client identification."""
@ -474,22 +490,24 @@ class TestRateLimitingMiddlewareIntegration:
rate_limit_server.add_middleware(
RateLimitingMiddleware(
max_requests_per_second=1.0, # Very slow refill to ensure rate limiting triggers
burst_capacity=4, # init + list_tools + call + list_tools = 4, so 2nd call fails
max_requests_per_second=0.001, # near-zero refill so tokens never replenish
burst_capacity=4,
get_client_id=get_client_id,
)
)
async with Client(rate_limit_server) as client:
# First request should succeed
await client.call_tool("quick_action", {"message": "first"})
# Second should be rate limited for this specific client
with pytest.raises(ToolError) as exc_info:
await client.call_tool("quick_action", {"message": "second"})
assert "Rate limit exceeded for client: test_client_123" in str(
exc_info.value
)
# Fire calls until the per-client bucket is exhausted; the exact
# burst count depends on the SDK's internal request volume.
exc_message = None
for i in range(10):
try:
await client.call_tool("quick_action", {"message": str(i)})
except MCPError as exc:
exc_message = str(exc)
break
assert exc_message is not None, "Rate limit was never triggered"
assert "Rate limit exceeded for client: test_client_123" in exc_message
async def test_global_rate_limiting(self, rate_limit_server):
"""Test global rate limiting across all clients."""
@ -505,7 +523,7 @@ class TestRateLimitingMiddlewareIntegration:
middleware.global_limiter.tokens = 0
with pytest.raises(ToolError, match="Global rate limit exceeded"):
with pytest.raises(MCPError, match="Global rate limit exceeded"):
await client.call_tool("quick_action", {"message": "blocked"})
async def test_rate_limiting_recovery_over_time(self, rate_limit_server):
@ -518,12 +536,16 @@ class TestRateLimitingMiddlewareIntegration:
)
async with Client(rate_limit_server) as client:
# Use up capacity
await client.call_tool("quick_action", {"message": "first"})
# Should be rate limited immediately
with pytest.raises(ToolError):
await client.call_tool("quick_action", {"message": "second"})
# Exhaust the burst; the exact number of internal requests before the
# limit trips varies across SDK versions, so fire until it blocks.
hit_limit = False
for i in range(10):
try:
await client.call_tool("quick_action", {"message": str(i)})
except MCPError:
hit_limit = True
break
assert hit_limit, "Rate limit was never triggered"
# Wait for token bucket to refill (150ms should be enough for ~1.5 tokens)
await asyncio.sleep(0.15)