Adding tests to #506

tests/conftest.py

@@ -203,7 +203,6 @@ def make_http_paths(netcdf_local_paths, request):
 def daily_xarray_dataset():
     return make_ds(nt=10)
 
-
 @pytest.fixture(scope="session")
 def daily_xarray_dataset_with_coordinateless_dimension(daily_xarray_dataset):
     """
@@ -215,6 +214,12 @@ def daily_xarray_dataset_with_coordinateless_dimension(daily_xarray_dataset):
     del ds["lon"]
     return ds
 
+@pytest.fixture(scope="session")
+def daily_xarray_dataset_with_extra_dimension_coordinates():
+    ds = make_ds(nt=11, add_extra_dim_coords=True)
+    ds['extra_dim_var'] = ds['extra_dim_coord']
+    return ds
+
 
 @pytest.fixture(scope="session")
 def netcdf_local_paths_sequential_1d(daily_xarray_dataset, tmpdir_factory):
@@ -303,6 +308,28 @@ def netcdf_local_paths_sequential_multivariable_with_coordinateless_dimension(
         file_type="netcdf4",
     )
 
+@pytest.fixture(scope='session')
+def netcdf_local_paths_sequential_with_extra_dimension_coordinate(
+    daily_xarray_dataset_with_extra_dimension_coordinates, tmpdir_factory
+    ):
+    return make_local_paths(
+        daily_xarray_dataset_with_extra_dimension_coordinates,
+        tmpdir_factory,
+        "D",
+        split_up_files_by_day,
+        file_type="netcdf4",
+    )
+
+@pytest.fixture(
+    scope="session",
+    params=[
+        lazy_fixture("netcdf_local_paths_sequential_with_extra_dimension_coordinate"),
+    ],
+)
+def netcdf_local_paths_sequential_extra_dimension_coordinate(request):
+    return request.param
+
+
 
 @pytest.fixture(
     scope="session",
@@ -380,7 +407,6 @@ def netcdf_local_paths_sequential_with_coordinateless_dimension(
 def netcdf_local_file_pattern_sequential(netcdf_local_paths_sequential):
     return make_file_pattern(netcdf_local_paths_sequential)
 
-
 @pytest.fixture(scope="session")
 def netcdf_local_file_pattern_sequential_multivariable(
     netcdf_local_paths_sequential_multivariable,
@@ -418,6 +444,12 @@ def netcdf_local_file_pattern_sequential_with_coordinateless_dimension(
     """
     return make_file_pattern(netcdf_local_paths_sequential_with_coordinateless_dimension)
 
+@pytest.fixture(scope='session')
+def netcdf_local_file_pattern_sequential_extra_dimension_coordinate(
+    netcdf_local_paths_sequential_extra_dimension_coordinate,
+    ):
+    return make_file_pattern(netcdf_local_paths_sequential_extra_dimension_coordinate)
+
 
 # Storage fixtures --------------------------------------------------------------------------------
 

tests/data_generation.py

@@ -3,12 +3,12 @@
 import xarray as xr
 
 
-def make_ds(nt=10, non_dim_coords=False):
+def make_ds(nt=10, non_dim_coords=False, add_extra_dim_coords=False):
     """Return a synthetic random xarray dataset."""
     np.random.seed(2)
     # TODO: change nt to 11 in order to catch the edge case where
     # items_per_input does not evenly divide the length of the sequence dimension
-    ny, nx = 18, 36
+    ny, nx, ne = 18, 36, 2
     time = pd.date_range(start="2010-01-01", periods=nt, freq="D")
     lon = (np.arange(nx) + 0.5) * 360 / nx
     lon_attrs = {"units": "degrees_east", "long_name": "longitude"}
@@ -28,6 +28,11 @@ def make_ds(nt=10, non_dim_coords=False):
     if non_dim_coords:
         coords["timestep"] = ("time", np.arange(nt))
         coords["baz"] = (("lat", "lon"), np.random.rand(ny, nx))
+
+    if add_extra_dim_coords:
+        # introduce a coordinate with a dimension not used in the data variables
+        coords["extra_dim_coord"] = (("extra_dim", "time"), np.random.rand(ne, nt))
+        coords["extra_dim"] = ("extra_dim", np.arange(ne))
 
     ds = xr.Dataset(
         {"bar": (dims, bar, bar_attrs), "foo": (dims, foo, foo_attrs)},

tests/test_end_to_end.py

@@ -19,7 +19,7 @@ def pipeline():
         yield p
 
 
-@pytest.mark.parametrize("target_chunks", [{"time": 1}, {"time": 2}, {"time": 3}])
+@pytest.mark.parametrize("target_chunks", [{"time": 1}, {"time": 2}, {"time": 3}, {'time':1, 'lon': 18}])
 def test_xarray_zarr(
     daily_xarray_dataset,
     netcdf_local_file_pattern_sequential,
@@ -45,7 +45,6 @@ def test_xarray_zarr(
     assert ds.time.encoding["chunks"] == (target_chunks["time"],)
     xr.testing.assert_equal(ds.load(), daily_xarray_dataset)
 
-
 def test_xarray_zarr_subpath(
     daily_xarray_dataset,
     netcdf_local_file_pattern_sequential,
@@ -67,3 +66,33 @@ def test_xarray_zarr_subpath(
 
     ds = xr.open_dataset(os.path.join(tmp_target_url, "subpath"), engine="zarr")
     xr.testing.assert_equal(ds.load(), daily_xarray_dataset)
+
+@pytest.mark.parametrize("target_chunks", [{"time": 1}, {"time": 2}, {"time": 3}])
+def test_xarray_zarr_extra_dimension_coordinate(
+    daily_xarray_dataset_with_extra_dimension_coordinates,
+    netcdf_local_file_pattern_sequential_extra_dimension_coordinate,
+    pipeline,
+    tmp_target_url,
+    target_chunks,
+):
+    # triggers https://github.com/pangeo-forge/pangeo-forge-recipes/issues/504
+    target_chunks['extra_dim'] = 2
+
+    pattern = netcdf_local_file_pattern_sequential_extra_dimension_coordinate
+
+    with pipeline as p:
+        (
+            p
+            | beam.Create(pattern.items())
+            | OpenWithXarray(file_type=pattern.file_type)
+            | StoreToZarr(
+                target_root=tmp_target_url,
+                store_name="store",
+                target_chunks=target_chunks,
+                combine_dims=pattern.combine_dim_keys,
+            )
+        )
+
+    ds = xr.open_dataset(os.path.join(tmp_target_url, "store"), engine="zarr")
+    assert ds.time.encoding["chunks"] == (target_chunks["time"],)
+    xr.testing.assert_equal(ds.load(), daily_xarray_dataset_with_extra_dimension_coordinates)