diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py b/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py index b477bf3f406..2ccba2b2055 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py +++ b/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(): diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_identity_env.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_identity_env.py index ac7082c2668..1c65adac4c6 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_identity_env.py +++ b/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(): diff --git a/tests/test_litellm/proxy/spend_tracking/test_savings.py b/tests/test_litellm/proxy/spend_tracking/test_savings.py index c3297ee6ae9..7dd18587df3 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_savings.py +++ b/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")