Spaces:
Running
Running
Jeremiah Lowin commited on
Commit ·
462d3af
1
Parent(s): 96b666d
Add proxy middleware test
Browse files
tests/server/middleware/test_middleware.py
CHANGED
|
@@ -75,9 +75,9 @@ class RecordingMiddleware(MCPMiddleware):
|
|
| 75 |
"""Assert that a hook was called a specific number of times."""
|
| 76 |
calls = self.get_calls(hook=hook, method=method)
|
| 77 |
actual_times = len(calls)
|
|
|
|
| 78 |
assert actual_times == times, (
|
| 79 |
-
f"Expected {
|
| 80 |
-
f"{f' for method {method!r}' if method else ''}, "
|
| 81 |
f"but was called {actual_times} times"
|
| 82 |
)
|
| 83 |
return True
|
|
@@ -546,3 +546,22 @@ class TestNestedMiddlewareHooks:
|
|
| 546 |
assert nested_middleware.assert_called(
|
| 547 |
hook="on_list_resource_templates", times=1
|
| 548 |
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 75 |
"""Assert that a hook was called a specific number of times."""
|
| 76 |
calls = self.get_calls(hook=hook, method=method)
|
| 77 |
actual_times = len(calls)
|
| 78 |
+
identifier = dict(hook=hook, method=method)
|
| 79 |
assert actual_times == times, (
|
| 80 |
+
f"Expected {times} calls for {identifier}, "
|
|
|
|
| 81 |
f"but was called {actual_times} times"
|
| 82 |
)
|
| 83 |
return True
|
|
|
|
| 546 |
assert nested_middleware.assert_called(
|
| 547 |
hook="on_list_resource_templates", times=1
|
| 548 |
)
|
| 549 |
+
|
| 550 |
+
|
| 551 |
+
class TestProxyServer:
|
| 552 |
+
async def test_call_tool(
|
| 553 |
+
self, mcp_server: FastMCP, recording_middleware: RecordingMiddleware
|
| 554 |
+
):
|
| 555 |
+
# proxy server will have its tools listed as well as called in order to
|
| 556 |
+
# run the `should_enable_component` hook prior to the call.
|
| 557 |
+
proxy_server = FastMCP.as_proxy(mcp_server, name="Proxy Server")
|
| 558 |
+
async with Client(proxy_server) as client:
|
| 559 |
+
await client.call_tool("add", {"a": 1, "b": 2})
|
| 560 |
+
|
| 561 |
+
assert recording_middleware.assert_called(times=6)
|
| 562 |
+
assert recording_middleware.assert_called(method="tools/call", times=3)
|
| 563 |
+
assert recording_middleware.assert_called(method="tools/list", times=3)
|
| 564 |
+
assert recording_middleware.assert_called(hook="on_message", times=2)
|
| 565 |
+
assert recording_middleware.assert_called(hook="on_request", times=2)
|
| 566 |
+
assert recording_middleware.assert_called(hook="on_call_tool", times=1)
|
| 567 |
+
assert recording_middleware.assert_called(hook="on_list_tools", times=1)
|