diff --git a/tests/utils/test_import_utils.py b/tests/utils/test_import_utils.py index 3d846174aca1..c5cde039b4ed 100644 --- a/tests/utils/test_import_utils.py +++ b/tests/utils/test_import_utils.py @@ -1,13 +1,16 @@ import sys - -from transformers.utils.import_utils import clear_import_cache - +import pytest +from transformers.utils.import_utils import clear_import_cache, _LazyModule def test_clear_import_cache(): # Import some transformers modules + from transformers import AutoModel, AutoTokenizer # Get initial module count - initial_modules = {name: mod for name, mod in sys.modules.items() if name.startswith("transformers.")} + initial_modules = { + name: mod for name, mod in sys.modules.items() + if name.startswith('transformers.') + } # Verify we have some modules loaded assert len(initial_modules) > 0 @@ -16,8 +19,12 @@ def test_clear_import_cache(): clear_import_cache() # Check modules were removed - remaining_modules = {name: mod for name, mod in sys.modules.items() if name.startswith("transformers.")} + remaining_modules = { + name: mod for name, mod in sys.modules.items() + if name.startswith('transformers.') + } assert len(remaining_modules) < len(initial_modules) # Verify we can reimport - assert "transformers.models.auto.modeling_auto" in sys.modules + from transformers import AutoModel, AutoTokenizer + assert 'transformers.models.auto.modeling_auto' in sys.modules \ No newline at end of file