test: isolate legacy DGL shim

tests/parity/conftest.py

@@ -68,6 +68,11 @@ def legacy_dataset_module_without_dgl():
     try:
         import dgl  # noqa: F401
     except ImportError:
+        missing = object()
+        originals = {
+            name: sys.modules.get(name, missing)
+            for name in ("dgl", "dgl.data", "matplotlib", "matplotlib.pyplot")
+        }
         dgl_module = types.ModuleType("dgl")
         dgl_data_module = types.ModuleType("dgl.data")
         dgl_data_module.DGLDataset = type("DGLDataset", (), {})
@@ -79,6 +84,15 @@ def legacy_dataset_module_without_dgl():
         matplotlib_module.pyplot = matplotlib_pyplot_module
         sys.modules["matplotlib"] = matplotlib_module
         sys.modules["matplotlib.pyplot"] = matplotlib_pyplot_module
+        try:
+            from root_gnn_base import dataset
+        finally:
+            for name, original in originals.items():
+                if original is missing:
+                    sys.modules.pop(name, None)
+                else:
+                    sys.modules[name] = original
+        return dataset
     from root_gnn_base import dataset
 
     return dataset