From b0adc8eb9c6a465651bd62425a9a890550758485 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Sun, 5 Jul 2026 23:03:47 -0400 Subject: [PATCH] Update middleware logging snapshots and rate-limit tests for SDK v2 (param shapes, MCPError surface) --- tests/server/middleware/test_logging.py | 21 ++-- tests/server/middleware/test_rate_limiting.py | 98 ++++++++++++------- 2 files changed, 70 insertions(+), 49 deletions(-) diff --git a/tests/server/middleware/test_logging.py b/tests/server/middleware/test_logging.py index 5c9360447..cb191dd16 100644 --- a/tests/server/middleware/test_logging.py +++ b/tests/server/middleware/test_logging.py @@ -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 """) diff --git a/tests/server/middleware/test_rate_limiting.py b/tests/server/middleware/test_rate_limiting.py index d2d81ef98..4ae5ae114 100644 --- a/tests/server/middleware/test_rate_limiting.py +++ b/tests/server/middleware/test_rate_limiting.py @@ -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)