Add tests for grib idx & reinflate

tests/test__grib_idx.py

@@ -16,6 +16,7 @@
 COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
 OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 """
+
 import copy
 import re
 import unittest
@@ -247,7 +248,13 @@ def test_build_idx_grib_mapping(self):
                         f"{mapping_fname}.idx_grib_mapping.parquet",
                     )
                     # Build the mapping from idx to cfgrib metadata and assert it matches the fixture
-                    pd.testing.assert_frame_equal(mapping, pd.read_parquet(test_path))
+                    expected = pd.read_parquet(test_path)
+                    pd.testing.assert_frame_equal(
+                        mapping,
+                        expected.assign(
+                            step=lambda x: x.step.astype("timedelta64[ns]")
+                        ),
+                    )
 
                     # parse the idx files for 20231104 and compare the mapped result to the direct indexed result
                     test_name = fnames["20231104"]
@@ -284,6 +291,10 @@ def test_build_idx_grib_mapping(self):
                     )
                     expected = pd.read_parquet(kindex_test_path)
 
+                    expected = expected.assign(
+                        step=lambda x: x.step.astype("timedelta64[ns]")
+                    )
+
                     expected = expected.set_index(
                         ["varname", "typeOfLevel", "stepType", "step", "level"]
                     ).sort_index()
@@ -382,7 +393,13 @@ def test_kerchunk_indexing(self):
                         TEST_DATE,
                         f"{fname}.kindex.parquet",
                     )
-                    pd.testing.assert_frame_equal(kindex, pd.read_parquet(test_path))
+                    expected = pd.read_parquet(test_path)
+                    pd.testing.assert_frame_equal(
+                        kindex,
+                        expected.assign(
+                            step=lambda x: x.step.astype("timedelta64[ns]")
+                        ),
+                    )
 
     @unittest.skip("TODO")
     def test_extract_dataset_chunk_index(self):
@@ -428,7 +445,9 @@ def _test_index_extraction(self, sample_prefix: str):
             THIS_DIR, "grib_idx_fixtures", sample_prefix, "kerchunk_index.parquet"
         )
         expected = pd.read_parquet(test_path)
-        pd.testing.assert_frame_equal(k_index, expected)
+        pd.testing.assert_frame_equal(
+            k_index, expected.assign(step=lambda x: x.step.astype("timedelta64[ns]"))
+        )
 
     def test_strip_datavar_chunks(self):
         for sample_prefix, pre, post in [
@@ -587,7 +606,7 @@ def _reinflate_grib_store(
             with self.subTest(node_path=node.path):
 
                 match aggregation:
-                    case (AggregationType.HORIZON | AggregationType.BEST_AVAILABLE):
+                    case AggregationType.HORIZON | AggregationType.BEST_AVAILABLE:
 
                         self.assertEqual(node.time.dims, node.valid_time.dims)
                         self.assertEqual(node.time.dims, node.step.dims)

tests/test_xarray_backend.py

@@ -12,6 +12,7 @@
     format="NETCDF3_CLASSIC"
 )
 
+
 # Add parameterize test for storage formats (.json, .parquet)
 def test_reference_netcdf(m):
     m.pipe("data.nc3", bdata)