From ce0323263eeb9d59ad8134ba7921ea53bffa074f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 22 Jun 2026 08:45:34 -0700 Subject: [PATCH] Fix test isolation: restore sys.modules after the pre-import gate test (#6578) * 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> --- studio/backend/tests/test_ssm_runtime.py | 46 +++++++++++++++++------- 1 file changed, 33 insertions(+), 13 deletions(-) diff --git a/studio/backend/tests/test_ssm_runtime.py b/studio/backend/tests/test_ssm_runtime.py index 2f6bae9b79..bb0caa2887 100644 --- a/studio/backend/tests/test_ssm_runtime.py +++ b/studio/backend/tests/test_ssm_runtime.py @@ -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():