Make create_test_datatree a pytest.fixture (#107)

* Migrate create_test_datatree to pytest.fixture

* Move create_test_datatree fixture to conftest.py

* black

* whatsnew

Co-authored-by: Thomas Nicholas <[email protected]>

Author: benjaminwoods

datatree/tests/conftest.py

@@ -0,0 +1,65 @@
+import pytest
+import xarray as xr
+
+from datatree import DataTree
+
+
+@pytest.fixture(scope="module")
+def create_test_datatree():
+    """
+    Create a test datatree with this structure:
+
+    <datatree.DataTree>
+    |-- set1
+    |   |-- <xarray.Dataset>
+    |   |   Dimensions:  ()
+    |   |   Data variables:
+    |   |       a        int64 0
+    |   |       b        int64 1
+    |   |-- set1
+    |   |-- set2
+    |-- set2
+    |   |-- <xarray.Dataset>
+    |   |   Dimensions:  (x: 2)
+    |   |   Data variables:
+    |   |       a        (x) int64 2, 3
+    |   |       b        (x) int64 0.1, 0.2
+    |   |-- set1
+    |-- set3
+    |-- <xarray.Dataset>
+    |   Dimensions:  (x: 2, y: 3)
+    |   Data variables:
+    |       a        (y) int64 6, 7, 8
+    |       set0     (x) int64 9, 10
+
+    The structure has deliberately repeated names of tags, variables, and
+    dimensions in order to better check for bugs caused by name conflicts.
+    """
+
+    def _create_test_datatree(modify=lambda ds: ds):
+        set1_data = modify(xr.Dataset({"a": 0, "b": 1}))
+        set2_data = modify(xr.Dataset({"a": ("x", [2, 3]), "b": ("x", [0.1, 0.2])}))
+        root_data = modify(xr.Dataset({"a": ("y", [6, 7, 8]), "set0": ("x", [9, 10])}))
+
+        # Avoid using __init__ so we can independently test it
+        root = DataTree(data=root_data)
+        set1 = DataTree(name="set1", parent=root, data=set1_data)
+        DataTree(name="set1", parent=set1)
+        DataTree(name="set2", parent=set1)
+        set2 = DataTree(name="set2", parent=root, data=set2_data)
+        DataTree(name="set1", parent=set2)
+        DataTree(name="set3", parent=root)
+
+        return root
+
+    return _create_test_datatree
+
+
+@pytest.fixture(scope="module")
+def simple_datatree(create_test_datatree):
+    """
+    Invoke create_test_datatree fixture (callback).
+
+    Returns a DataTree.
+    """
+    return create_test_datatree()

datatree/tests/test_dataset_api.py

@@ -4,8 +4,6 @@
 from datatree import DataTree
 from datatree.testing import assert_equal
 
-from .test_datatree import create_test_datatree
-
 
 class TestDSMethodInheritance:
     def test_dataset_method(self):
@@ -93,7 +91,7 @@ def test_binary_op_on_datatree(self):
 
 
 class TestUFuncs:
-    def test_tree(self):
+    def test_tree(self, create_test_datatree):
         dt = create_test_datatree()
         expected = create_test_datatree(modify=lambda ds: np.sin(ds))
         result_tree = np.sin(dt)

datatree/tests/test_datatree.py

@@ -5,52 +5,6 @@
 from datatree import DataTree
 
 
-def create_test_datatree(modify=lambda ds: ds):
-    """
-    Create a test datatree with this structure:
-
-    <datatree.DataTree>
-    |-- set1
-    |   |-- <xarray.Dataset>
-    |   |   Dimensions:  ()
-    |   |   Data variables:
-    |   |       a        int64 0
-    |   |       b        int64 1
-    |   |-- set1
-    |   |-- set2
-    |-- set2
-    |   |-- <xarray.Dataset>
-    |   |   Dimensions:  (x: 2)
-    |   |   Data variables:
-    |   |       a        (x) int64 2, 3
-    |   |       b        (x) int64 0.1, 0.2
-    |   |-- set1
-    |-- set3
-    |-- <xarray.Dataset>
-    |   Dimensions:  (x: 2, y: 3)
-    |   Data variables:
-    |       a        (y) int64 6, 7, 8
-    |       set0     (x) int64 9, 10
-
-    The structure has deliberately repeated names of tags, variables, and
-    dimensions in order to better check for bugs caused by name conflicts.
-    """
-    set1_data = modify(xr.Dataset({"a": 0, "b": 1}))
-    set2_data = modify(xr.Dataset({"a": ("x", [2, 3]), "b": ("x", [0.1, 0.2])}))
-    root_data = modify(xr.Dataset({"a": ("y", [6, 7, 8]), "set0": ("x", [9, 10])}))
-
-    # Avoid using __init__ so we can independently test it
-    root = DataTree(data=root_data)
-    set1 = DataTree(name="set1", parent=root, data=set1_data)
-    DataTree(name="set1", parent=set1)
-    DataTree(name="set2", parent=set1)
-    set2 = DataTree(name="set2", parent=root, data=set2_data)
-    DataTree(name="set1", parent=set2)
-    DataTree(name="set3", parent=root)
-
-    return root
-
-
 class TestTreeCreation:
     def test_empty(self):
         dt = DataTree(name="root")
@@ -322,8 +276,8 @@ def test_nones(self):
         assert [node.path for node in dt.subtree] == ["/", "/d", "/d/e"]
         xrt.assert_equal(dt["d/e"].ds, xr.Dataset())
 
-    def test_full(self):
-        dt = create_test_datatree()
+    def test_full(self, simple_datatree):
+        dt = simple_datatree
         paths = list(node.path for node in dt.subtree)
         assert paths == [
             "/",
@@ -335,16 +289,16 @@ def test_full(self):
             "/set3",
         ]
 
-    def test_roundtrip(self):
-        dt = create_test_datatree()
+    def test_roundtrip(self, simple_datatree):
+        dt = simple_datatree
         roundtrip = DataTree.from_dict(dt.to_dict())
         assert roundtrip.equals(dt)
 
     @pytest.mark.xfail
-    def test_roundtrip_unnamed_root(self):
+    def test_roundtrip_unnamed_root(self, simple_datatree):
         # See GH81
 
-        dt = create_test_datatree()
+        dt = simple_datatree
         dt.name = "root"
         roundtrip = DataTree.from_dict(dt.to_dict())
         assert roundtrip.equals(dt)

datatree/tests/test_formatting.py

@@ -5,8 +5,6 @@
 from datatree import DataTree
 from datatree.formatting import diff_tree_repr
 
-from .test_datatree import create_test_datatree
-
 
 class TestRepr:
     def test_print_empty_node(self):
@@ -50,8 +48,8 @@ def test_nested_node(self):
         printout = root.__str__()
         assert printout.splitlines()[2].startswith("    ")
 
-    def test_print_datatree(self):
-        dt = create_test_datatree()
+    def test_print_datatree(self, simple_datatree):
+        dt = simple_datatree
         print(dt)
 
         # TODO work out how to test something complex like this

datatree/tests/test_io.py

@@ -3,27 +3,26 @@
 from datatree.io import open_datatree
 from datatree.testing import assert_equal
 from datatree.tests import requires_h5netcdf, requires_netCDF4, requires_zarr
-from datatree.tests.test_datatree import create_test_datatree
 
 
 class TestIO:
     @requires_netCDF4
-    def test_to_netcdf(self, tmpdir):
+    def test_to_netcdf(self, tmpdir, simple_datatree):
         filepath = str(
             tmpdir / "test.nc"
         )  # casting to str avoids a pathlib bug in xarray
-        original_dt = create_test_datatree()
+        original_dt = simple_datatree
         original_dt.to_netcdf(filepath, engine="netcdf4")
 
         roundtrip_dt = open_datatree(filepath)
         assert_equal(original_dt, roundtrip_dt)
 
     @requires_netCDF4
-    def test_netcdf_encoding(self, tmpdir):
+    def test_netcdf_encoding(self, tmpdir, simple_datatree):
         filepath = str(
             tmpdir / "test.nc"
         )  # casting to str avoids a pathlib bug in xarray
-        original_dt = create_test_datatree()
+        original_dt = simple_datatree
 
         # add compression
         comp = dict(zlib=True, complevel=9)
@@ -40,35 +39,35 @@ def test_netcdf_encoding(self, tmpdir):
             original_dt.to_netcdf(filepath, encoding=enc, engine="netcdf4")
 
     @requires_h5netcdf
-    def test_to_h5netcdf(self, tmpdir):
+    def test_to_h5netcdf(self, tmpdir, simple_datatree):
         filepath = str(
             tmpdir / "test.nc"
         )  # casting to str avoids a pathlib bug in xarray
-        original_dt = create_test_datatree()
+        original_dt = simple_datatree
         original_dt.to_netcdf(filepath, engine="h5netcdf")
 
         roundtrip_dt = open_datatree(filepath)
         assert_equal(original_dt, roundtrip_dt)
 
     @requires_zarr
-    def test_to_zarr(self, tmpdir):
+    def test_to_zarr(self, tmpdir, simple_datatree):
         filepath = str(
             tmpdir / "test.zarr"
         )  # casting to str avoids a pathlib bug in xarray
-        original_dt = create_test_datatree()
+        original_dt = simple_datatree
         original_dt.to_zarr(filepath)
 
         roundtrip_dt = open_datatree(filepath, engine="zarr")
         assert_equal(original_dt, roundtrip_dt)
 
     @requires_zarr
-    def test_zarr_encoding(self, tmpdir):
+    def test_zarr_encoding(self, tmpdir, simple_datatree):
         import zarr
 
         filepath = str(
             tmpdir / "test.zarr"
         )  # casting to str avoids a pathlib bug in xarray
-        original_dt = create_test_datatree()
+        original_dt = simple_datatree
 
         comp = {"compressor": zarr.Blosc(cname="zstd", clevel=3, shuffle=2)}
         enc = {"/set2": {var: comp for var in original_dt["/set2"].ds.data_vars}}
@@ -83,26 +82,26 @@ def test_zarr_encoding(self, tmpdir):
             original_dt.to_zarr(filepath, encoding=enc, engine="zarr")
 
     @requires_zarr
-    def test_to_zarr_zip_store(self, tmpdir):
+    def test_to_zarr_zip_store(self, tmpdir, simple_datatree):
         from zarr.storage import ZipStore
 
         filepath = str(
             tmpdir / "test.zarr.zip"
         )  # casting to str avoids a pathlib bug in xarray
-        original_dt = create_test_datatree()
+        original_dt = simple_datatree
         store = ZipStore(filepath)
         original_dt.to_zarr(store)
 
         roundtrip_dt = open_datatree(store, engine="zarr")
         assert_equal(original_dt, roundtrip_dt)
 
     @requires_zarr
-    def test_to_zarr_not_consolidated(self, tmpdir):
+    def test_to_zarr_not_consolidated(self, tmpdir, simple_datatree):
         filepath = tmpdir / "test.zarr"
         zmetadata = filepath / ".zmetadata"
         s1zmetadata = filepath / "set1" / ".zmetadata"
         filepath = str(filepath)  # casting to str avoids a pathlib bug in xarray
-        original_dt = create_test_datatree()
+        original_dt = simple_datatree
         original_dt.to_zarr(filepath, consolidated=False)
         assert not zmetadata.exists()
         assert not s1zmetadata.exists()