scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

fix(tests): restore module attributes after reload in mcp identity env tests

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Devin AI 2026-09-01 10:59:08 +00:00
parent 2019c309ea
commit 95c7ca8801
1 changed files with 14 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_identity_env.py
View file

@ -23,6 +23,14 @@ 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 the pre-reload module attributes afterwards instead of reloading
# a third time: a reload re-creates every class in the module, so modules
# that imported names like MCPMissingUserEnvVarsError before this test
# would keep raising the old class while pytest.raises in later tests
# matches the new one
snapshots = {module: dict(vars(module)) for module in (utils_module, mgmt_module)}
def _apply_env(values):
for key, value in values.items():
@ -32,8 +40,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 +49,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():