force step dtypes in expected values

fsspec · Dec 16, 2024 · b6cc52c · b6cc52c
1 parent 261ee2f
commit b6cc52c
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/tests/test__grib_idx.py b/tests/test__grib_idx.py
@@ -247,7 +247,8 @@ 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 +285,8 @@ 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 +385,8 @@ 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 +432,7 @@ 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 [