Merge pull request #38891 from BerriAI/litellm_deflake_20260830

test: deflake MCP registry state, savings cost map, and MCP identity env reload tests

tests/test_litellm/proxy/_experimental/mcp_server/conftest.py

@@ -2,6 +2,30 @@ import os
 
 import pytest
 
+from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
+    global_mcp_server_manager,
+)
+
+
+@pytest.fixture(autouse=True)
+def _hermetic_mcp_server_registry():
+    """Restore the singleton ``global_mcp_server_manager``'s registry state around every
+    test, so entries seeded by one test never leak into another on a shared shard."""
+    saved_registry = dict(global_mcp_server_manager.registry)
+    saved_config_servers = dict(global_mcp_server_manager.config_mcp_servers)
+    saved_tool_mapping = dict(global_mcp_server_manager.tool_name_to_mcp_server_name_mapping)
+    saved_oauth_slots = global_mcp_server_manager._oauth_discovery_slots
+    try:
+        yield
+    finally:
+        global_mcp_server_manager.registry.clear()
+        global_mcp_server_manager.registry.update(saved_registry)
+        global_mcp_server_manager.config_mcp_servers.clear()
+        global_mcp_server_manager.config_mcp_servers.update(saved_config_servers)
+        global_mcp_server_manager.tool_name_to_mcp_server_name_mapping.clear()
+        global_mcp_server_manager.tool_name_to_mcp_server_name_mapping.update(saved_tool_mapping)
+        global_mcp_server_manager._oauth_discovery_slots = saved_oauth_slots
+
 
 @pytest.fixture(autouse=True)
 def _hermetic_server_root_path():

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_identity_env.py

@@ -23,6 +23,12 @@ MGMT_MODULE = "litellm.proxy.management_endpoints.mcp_management_endpoints"
 @contextlib.contextmanager
 def _env_and_reload(**env):
     saved = {key: os.environ.get(key) for key in env}
+    utils_module = importlib.import_module(UTILS_MODULE)
+    mgmt_module = importlib.import_module(MGMT_MODULE)
+    # Restore pre-reload module attributes afterwards instead of reloading again:
+    # a reload re-creates the module's classes, breaking exception identity for
+    # modules that imported them earlier
+    snapshots = {module: dict(vars(module)) for module in (utils_module, mgmt_module)}
 
     def _apply_env(values):
         for key, value in values.items():
@@ -32,8 +38,8 @@ def _env_and_reload(**env):
                 os.environ[key] = value
 
     def _reload():
-        utils = importlib.reload(importlib.import_module(UTILS_MODULE))
-        mgmt = importlib.reload(importlib.import_module(MGMT_MODULE))
+        utils = importlib.reload(utils_module)
+        mgmt = importlib.reload(mgmt_module)
         return utils, mgmt
 
     try:
@@ -41,7 +47,10 @@ def _env_and_reload(**env):
         yield _reload()
     finally:
         _apply_env(saved)
-        _reload()
+        for module, snapshot in snapshots.items():
+            for key in [key for key in vars(module) if key not in snapshot]:
+                delattr(module, key)
+            vars(module).update(snapshot)
 
 
 def test_defaults_used_when_env_unset():

tests/test_litellm/proxy/spend_tracking/test_savings.py

@@ -14,6 +14,8 @@ from litellm.proxy.spend_tracking.savings import (
 from litellm.router import Router
 from litellm.types.utils import Usage
 
+pytestmark = pytest.mark.usefixtures("local_model_cost_map")
+
 
 def _anthropic_costs(model: str) -> tuple[float, float]:
     info = litellm.get_model_info(model=model, custom_llm_provider="anthropic")