update

Author: justusschock

tests/tests_fabric/graveyard/test_tpu.py

@@ -18,6 +18,10 @@ def test_graveyard_single_tpu(import_path, name, tpu_available):
     with pytest.deprecated_call(match="is deprecated"):
         cls(device)
 
+    # required to prevent env-var leakage
+    if hasattr(cls, "teardown"):
+        cls.teardown()
+
 
 @pytest.mark.parametrize(
     ("import_path", "name"),
@@ -39,3 +43,7 @@ def test_graveyard_no_device(import_path, name, tpu_available):
     cls = getattr(module, name)
     with pytest.deprecated_call(match="is deprecated"):
         cls()
+
+    # required to prevent env-var leakage
+    if hasattr(cls, "teardown"):
+        cls.teardown()

tests/tests_pytorch/graveyard/test_tpu.py

@@ -18,6 +18,10 @@ def test_graveyard_single_tpu(import_path, name):
     with pytest.deprecated_call(match="is deprecated"), pytest.raises(ModuleNotFoundError, match="torch_xla"):
         cls(device)
 
+    # required to prevent env-var leakage
+    if hasattr(cls, "teardown"):
+        cls.teardown()
+
 
 @pytest.mark.parametrize(
     ("import_path", "name"),
@@ -39,3 +43,7 @@ def test_graveyard_no_device(import_path, name):
     cls = getattr(module, name)
     with pytest.deprecated_call(match="is deprecated"), pytest.raises(ModuleNotFoundError, match="torch_xla"):
         cls()
+
+    # required to prevent env-var leakage
+    if hasattr(cls, "teardown"):
+        cls.teardown()

tests/tests_pytorch/utilities/migration/test_utils.py

@@ -125,7 +125,10 @@ def test_patch_legacy_imports_unified(pl_version):
     assert any(key.startswith("lightning." + "pytorch") for key in sys.modules), (
         f"Imported unified package, so it has to be in sys.modules: {_list_sys_modules('lightning' + '.pytorch')}"
     )
-    assert not any(key.startswith("pytorch_lightning") for key in sys.modules), (
+    assert not any(
+        key.startswith("pytorch_lightning") and not key.startswith("pytorch_lightning_enterprise")
+        for key in sys.modules
+    ), (
         "Should not import standalone package, all imports should be redirected to the unified package;\n"
         f" environment: {_list_sys_modules('pytorch_lightning')}"
     )