Spaces:
Running
Running
File size: 11,310 Bytes
5800457 5bfd7a8 5800457 5bfd7a8 5800457 5bfd7a8 5800457 5bfd7a8 5800457 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 | """Tests for timing middleware."""
import asyncio
import logging
import time
from unittest.mock import AsyncMock, MagicMock
import pytest
from fastmcp import FastMCP
from fastmcp.client import Client
from fastmcp.server.middleware.middleware import MiddlewareContext
from fastmcp.server.middleware.timing import DetailedTimingMiddleware, TimingMiddleware
@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 TestTimingMiddleware:
"""Test timing middleware functionality."""
def test_init_default(self):
"""Test default initialization."""
middleware = TimingMiddleware()
assert middleware.logger.name == "fastmcp.timing"
assert middleware.log_level == logging.INFO
def test_init_custom(self):
"""Test custom initialization."""
logger = logging.getLogger("custom")
middleware = TimingMiddleware(logger=logger, log_level=logging.DEBUG)
assert middleware.logger is logger
assert middleware.log_level == logging.DEBUG
async def test_on_request_success(self, mock_context, mock_call_next, caplog):
"""Test timing successful requests."""
middleware = TimingMiddleware()
with caplog.at_level(logging.INFO):
result = await middleware.on_request(mock_context, mock_call_next)
assert result == "test_result"
assert mock_call_next.called
assert "Request test_method completed in" in caplog.text
assert "ms" in caplog.text
async def test_on_request_failure(self, mock_context, caplog):
"""Test timing failed requests."""
middleware = TimingMiddleware()
mock_call_next = AsyncMock(side_effect=ValueError("test error"))
with caplog.at_level(logging.INFO):
with pytest.raises(ValueError):
await middleware.on_request(mock_context, mock_call_next)
assert "Request test_method failed after" in caplog.text
assert "ms: test error" in caplog.text
class TestDetailedTimingMiddleware:
"""Test detailed timing middleware functionality."""
def test_init_default(self):
"""Test default initialization."""
middleware = DetailedTimingMiddleware()
assert middleware.logger.name == "fastmcp.timing.detailed"
assert middleware.log_level == logging.INFO
async def test_on_call_tool(self, caplog):
"""Test timing tool calls."""
middleware = DetailedTimingMiddleware()
context = MagicMock()
context.message.name = "test_tool"
mock_call_next = AsyncMock(return_value="tool_result")
with caplog.at_level(logging.INFO):
result = await middleware.on_call_tool(context, mock_call_next)
assert result == "tool_result"
assert "Tool 'test_tool' completed in" in caplog.text
async def test_on_read_resource(self, caplog):
"""Test timing resource reads."""
middleware = DetailedTimingMiddleware()
context = MagicMock()
context.message.uri = "test://resource"
mock_call_next = AsyncMock(return_value="resource_result")
with caplog.at_level(logging.INFO):
result = await middleware.on_read_resource(context, mock_call_next)
assert result == "resource_result"
assert "Resource 'test://resource' completed in" in caplog.text
async def test_on_get_prompt(self, caplog):
"""Test timing prompt retrieval."""
middleware = DetailedTimingMiddleware()
context = MagicMock()
context.message.name = "test_prompt"
mock_call_next = AsyncMock(return_value="prompt_result")
with caplog.at_level(logging.INFO):
result = await middleware.on_get_prompt(context, mock_call_next)
assert result == "prompt_result"
assert "Prompt 'test_prompt' completed in" in caplog.text
async def test_on_list_tools(self, caplog):
"""Test timing tool listing."""
middleware = DetailedTimingMiddleware()
context = MagicMock()
mock_call_next = AsyncMock(return_value="tools_result")
with caplog.at_level(logging.INFO):
result = await middleware.on_list_tools(context, mock_call_next)
assert result == "tools_result"
assert "List tools completed in" in caplog.text
async def test_operation_failure(self, caplog):
"""Test timing failed operations."""
middleware = DetailedTimingMiddleware()
context = MagicMock()
context.message.name = "failing_tool"
mock_call_next = AsyncMock(side_effect=RuntimeError("operation failed"))
with caplog.at_level(logging.INFO):
with pytest.raises(RuntimeError):
await middleware.on_call_tool(context, mock_call_next)
assert "Tool 'failing_tool' failed after" in caplog.text
assert "ms: operation failed" in caplog.text
@pytest.fixture
def timing_server():
"""Create a FastMCP server specifically for timing middleware tests."""
mcp = FastMCP("TimingTestServer")
@mcp.tool
def instant_task() -> str:
"""A task that completes instantly."""
return "Done instantly"
@mcp.tool
def short_task() -> str:
"""A task that takes 0.1 seconds."""
time.sleep(0.1)
return "Done after 0.1s"
@mcp.tool
def medium_task() -> str:
"""A task that takes 0.15 seconds."""
time.sleep(0.15)
return "Done after 0.15s"
@mcp.tool
def failing_task() -> str:
"""A task that always fails."""
raise ValueError("Task failed as expected")
@mcp.resource("timer://test")
def test_resource() -> str:
"""A resource that takes time to read."""
time.sleep(0.05)
return "Resource content after 0.05s"
@mcp.prompt
def test_prompt() -> str:
"""A prompt that takes time to generate."""
time.sleep(0.08)
return "Prompt content after 0.08s"
return mcp
class TestTimingMiddlewareIntegration:
"""Integration tests for timing middleware with real FastMCP server."""
async def test_timing_middleware_measures_tool_execution(
self, timing_server, caplog
):
"""Test that timing middleware accurately measures tool execution times."""
timing_server.add_middleware(TimingMiddleware())
with caplog.at_level(logging.INFO):
async with Client(timing_server) as client:
# Test instant task
await client.call_tool("instant_task")
# Test short task (0.1s)
await client.call_tool("short_task")
# Test medium task (0.15s)
await client.call_tool("medium_task")
log_text = caplog.text
# Should have timing logs for all three calls (plus any extra list_tools calls)
timing_logs = [
line
for line in log_text.split("\n")
if "completed in" in line and "ms" in line
]
assert (
len(timing_logs) >= 3
) # At least 3 tool calls, may have additional list_tools calls
# Verify that longer tasks show longer timing (roughly)
assert "tools/call completed in" in log_text
assert "ms" in log_text
async def test_timing_middleware_handles_failures(self, timing_server, caplog):
"""Test that timing middleware measures time even for failed operations."""
timing_server.add_middleware(TimingMiddleware())
with caplog.at_level(logging.INFO):
async with Client(timing_server) as client:
# This should fail but still be timed
with pytest.raises(Exception):
await client.call_tool("failing_task")
# Should log the failure with timing
assert "tools/call failed after" in caplog.text
assert "ms:" in caplog.text
async def test_detailed_timing_middleware_per_operation(
self, timing_server, caplog
):
"""Test that detailed timing middleware provides operation-specific timing."""
timing_server.add_middleware(DetailedTimingMiddleware())
with caplog.at_level(logging.INFO):
async with Client(timing_server) as client:
# Test tool call
await client.call_tool("short_task")
# Test resource read
await client.read_resource("timer://test")
# Test prompt
await client.get_prompt("test_prompt")
# Test listing operations
await client.list_tools()
await client.list_resources()
await client.list_prompts()
log_text = caplog.text
# Should have specific timing logs for each operation type
assert "Tool 'short_task' completed in" in log_text
assert "Resource 'timer://test' completed in" in log_text
assert "Prompt 'test_prompt' completed in" in log_text
assert "List tools completed in" in log_text
assert "List resources completed in" in log_text
assert "List prompts completed in" in log_text
async def test_timing_middleware_concurrent_operations(self, timing_server, caplog):
"""Test timing middleware with concurrent operations."""
timing_server.add_middleware(TimingMiddleware())
with caplog.at_level(logging.INFO):
async with Client(timing_server) as client:
# Run multiple operations concurrently
tasks = [
client.call_tool("instant_task"),
client.call_tool("short_task"),
client.call_tool("instant_task"),
]
await asyncio.gather(*tasks)
log_text = caplog.text
# Should have timing logs for all concurrent operations (including extra list_tools calls)
timing_logs = [line for line in log_text.split("\n") if "completed in" in line]
assert (
len(timing_logs) >= 3
) # At least 3 tool calls, may have additional list_tools calls
async def test_timing_middleware_custom_logger(self, timing_server):
"""Test timing middleware with custom logger configuration."""
import io
import logging
# Create a custom logger that writes to a string buffer
log_buffer = io.StringIO()
handler = logging.StreamHandler(log_buffer)
custom_logger = logging.getLogger("custom_timing")
custom_logger.addHandler(handler)
custom_logger.setLevel(logging.DEBUG)
# Use custom logger and log level
timing_server.add_middleware(
TimingMiddleware(logger=custom_logger, log_level=logging.DEBUG)
)
async with Client(timing_server) as client:
await client.call_tool("instant_task")
# Check that our custom logger was used
log_output = log_buffer.getvalue()
assert "tools/call completed in" in log_output
assert "ms" in log_output
|