Add vars and dims tests

Author: jacobbieker

ocf_datapipes/validation/check_vars_and_dims.py

@@ -13,8 +13,8 @@ class CheckVarsAndDimsIterDataPipe(IterDataPipe):
     def __init__(
         self,
         source_datapipe: IterDataPipe,
-        expected_dimensions: Iterable[str],
-        expected_data_vars: Iterable[str],
+        expected_dimensions: Optional[Iterable[str]] = None,
+        expected_data_vars: Optional[Iterable[str]] = None,
         dataset_name: Optional[str] = None,
     ):
         """
@@ -41,19 +41,23 @@ def __iter__(self) -> Union[xr.DataArray, xr.Dataset]:
         """
         for xr_data in self.source_datapipe:
             if self.dataset_name is None:
-                xr_data = validate_data_vars(xr_data, self.expected_data_vars)
-                xr_data = validate_dims(xr_data, self.expected_dimensions)
-                xr_data = validate_coords(xr_data, self.expected_dimensions)
+                if self.expected_data_vars is not None:
+                    xr_data = validate_data_vars(xr_data, self.expected_data_vars)
+                if self.expected_dimensions is not None:
+                    xr_data = validate_dims(xr_data, self.expected_dimensions)
+                    xr_data = validate_coords(xr_data, self.expected_dimensions)
             else:
-                xr_data[self.dataset_name] = validate_data_vars(
-                    xr_data[self.dataset_name], self.expected_data_vars
-                )
-                xr_data[self.dataset_name] = validate_dims(
-                    xr_data[self.dataset_name], self.expected_dimensions
-                )
-                xr_data[self.dataset_name] = validate_coords(
-                    xr_data[self.dataset_name], self.expected_dimensions
-                )
+                if self.expected_data_vars is not None:
+                    xr_data[self.dataset_name] = validate_data_vars(
+                        xr_data[self.dataset_name], self.expected_data_vars
+                    )
+                if self.expected_dimensions is not None:
+                    xr_data[self.dataset_name] = validate_dims(
+                        xr_data[self.dataset_name], self.expected_dimensions
+                    )
+                    xr_data[self.dataset_name] = validate_coords(
+                        xr_data[self.dataset_name], self.expected_dimensions
+                    )
             yield xr_data
 
 

tests/validation/test_check_vars_and_dims.py

@@ -1,5 +1,25 @@
 from ocf_datapipes.validation import CheckVarsAndDims
 
 
-def test_check_vars_and_dims():
-    pass
+def test_check_vars_and_dims_gsp(gsp_dp):
+    gsp_dp = gsp_dp.check_vars_and_dims(expected_dimensions=("time_utc", "gsp_id"))
+    next(iter(gsp_dp))
+
+def test_check_vars_and_dims_sat(sat_dp):
+    sat_dp = sat_dp.check_vars_and_dims(expected_dimensions=("time_utc", "channel", "x_geostationary", "y_geostationary"))
+    next(iter(sat_dp))
+
+def test_check_vars_and_dim_passiv(passiv_dp):
+    passiv_dp = passiv_dp.check_vars_and_dims(
+        expected_dimensions=("time_utc", "pv_system_id"))
+    next(iter(passiv_dp))
+
+def test_check_vars_and_dim_nwp(nwp_dp):
+    nwp_dp = nwp_dp.check_vars_and_dims(
+        expected_dimensions=("init_time_utc", "channel", "step", "x_osgb", "y_osgb"))
+    next(iter(nwp_dp))
+
+def test_check_vars_and_dim_topo(topo_dp):
+    topo_dp = topo_dp.check_vars_and_dims(
+        expected_dimensions=("x_osgb", "y_osgb"))
+    next(iter(topo_dp))