From 205b08703bd94430587352da0d3b793ddd440183 Mon Sep 17 00:00:00 2001 From: Alexsander Hamir Date: Tue, 10 Feb 2026 10:26:44 -0800 Subject: [PATCH] fix: resolve test failures - budget mock_table and MCP openapi async - test_budget_endpoints: define mock_table in client_and_mocks fixture to fix NameError in 5 tests - mcp_server_manager: make _register_openapi_tools async and use load_openapi_spec_async to avoid RuntimeError when called from running event loop; await in load_servers_from_config - test_mcp_server_manager: await _register_openapi_tools in test_register_openapi_tools_includes_static_headers --- .../_experimental/mcp_server/mcp_server_manager.py | 12 +++++++----- .../mcp_server/test_mcp_server_manager.py | 2 +- .../management_endpoints/test_budget_endpoints.py | 1 + 3 files changed, 9 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 532aea249bf..24eae430d85 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -341,7 +341,7 @@ class MCPServerManager: verbose_logger.info( f"Loading OpenAPI spec from {spec_path} for server {server_name}" ) - self._register_openapi_tools( + await self._register_openapi_tools( spec_path=spec_path, server=new_server, base_url=server_config.get("url", ""), @@ -353,7 +353,9 @@ class MCPServerManager: self.initialize_tool_name_to_mcp_server_name_mapping() - def _register_openapi_tools(self, spec_path: str, server: MCPServer, base_url: str): + async def _register_openapi_tools( + self, spec_path: str, server: MCPServer, base_url: str + ): """ Register tools from an OpenAPI specification for a given server. @@ -375,15 +377,15 @@ class MCPServerManager: get_base_url as get_openapi_base_url, ) from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( - load_openapi_spec, + load_openapi_spec_async, ) from litellm.proxy._experimental.mcp_server.tool_registry import ( global_mcp_tool_registry, ) try: - # Load OpenAPI spec - spec = load_openapi_spec(spec_path) + # Load OpenAPI spec (async to avoid "called from within a running event loop") + spec = await load_openapi_spec_async(spec_path) # Use base_url from config if provided, otherwise extract from spec if not base_url: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index abb8dd49159..b4b5811666b 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -1043,7 +1043,7 @@ class TestMCPServerManager: "litellm.proxy._experimental.mcp_server.tool_registry.global_mcp_tool_registry.register_tool", return_value=None, ): - manager._register_openapi_tools( + await manager._register_openapi_tools( spec_path=str(spec_path), server=server, base_url="https://example.com", diff --git a/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py index d8c505223d9..b15b9d622e4 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py @@ -21,6 +21,7 @@ sys.path.insert( def client_and_mocks(monkeypatch): # Setup MagicMock Prisma mock_prisma = MagicMock() + mock_table = MagicMock() mock_table.create = AsyncMock(side_effect=lambda *, data: data) mock_table.update = AsyncMock(side_effect=lambda *, where, data: {**where, **data})