ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix test isolation: restore sys.modules after the pre-import gate test (#6578)

Browse source 
* Restore sys.modules in test_pre_import_gate_is_transformers_free

The test pops transformers and utils.models.model_config from sys.modules to
assert the pre-import security gate does not re-import them, but never put them
back. A later importer then rebound a fresh utils.models.model_config, so tests
that had captured the original instance missed their patches and hit the real
path: test_vision_cache patches _is_vision_model_uncached on the original
module, but is_vision_model (still bound to that original) ran the real network
lookup instead. This produced 17 spurious failures whenever test_ssm_runtime
ran before test_vision_cache in the same process.

Snapshot the removed modules and restore the original objects in a finally, so
the assertions still run against a clean slate while later tests see the same
module instances they captured at import time.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-06-22 08:45:34 -07:00 • committed by GitHub
commit ce0323263e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 33 additions and 13 deletions
Download patch file Download diff file Expand all files Collapse all files

46
studio/backend/tests/test_ssm_runtime.py
View file

@ -445,20 +445,40 @@ def test_pre_import_gate_is_transformers_free():
import utils.security.file_security as fs
import utils.security.consent as consent
for m in list(_sys.modules):
if m == "transformers" or m.startswith("transformers.") or m == "utils.models.model_config":
def _is_gated_module(name: str) -> bool:
return (
name == "transformers"
or name.startswith("transformers.")
or name == "utils.models.model_config"
)
# Snapshot then remove the modules so we can assert the gate does not re-import them.
# Restore the originals afterwards (finally): popping utils.models.model_config without
# restoring it makes a later importer get a fresh instance, so tests that patched the
# first instance (e.g. test_vision_cache) miss and hit the real network path.
_saved = {m: _sys.modules[m] for m in list(_sys.modules) if _is_gated_module(m)}
for m in _saved:
_sys.modules.pop(m, None)
try:
with patch.object(fs, "_fetch_security_status", return_value = None):
fs.evaluate_file_security("nvidia/Nemotron-H-8B", load_subdirs = ())
with patch.object(
consent, "_load_remote_code_configs", return_value = [{"model_type": "nemotron_h"}]
):
from utils.security import evaluate_remote_code_consent_for_targets
evaluate_remote_code_consent_for_targets(
["nvidia/Nemotron-H-8B"], trust_remote_code = True
)
assert "transformers" not in _sys.modules
assert "utils.models.model_config" not in _sys.modules
finally:
# Drop anything the gate imported, then rebind the original module objects so later
# tests see the same instances they captured at import time.
for m in [m for m in list(_sys.modules) if _is_gated_module(m) and m not in _saved]:
_sys.modules.pop(m, None)
with patch.object(fs, "_fetch_security_status", return_value = None):
fs.evaluate_file_security("nvidia/Nemotron-H-8B", load_subdirs = ())
with patch.object(
consent, "_load_remote_code_configs", return_value = [{"model_type": "nemotron_h"}]
):
from utils.security import evaluate_remote_code_consent_for_targets
evaluate_remote_code_consent_for_targets(["nvidia/Nemotron-H-8B"], trust_remote_code = True)
assert "transformers" not in _sys.modules
assert "utils.models.model_config" not in _sys.modules
_sys.modules.update(_saved)
def test_pre_import_gate_skips_subdir_computation():