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 {hook!r} to be called {times} times"
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)