Fix tests and black format

fsspec · Dec 16, 2024 · 72f684e · 72f684e
1 parent 261ee2f
commit 72f684e
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 4 deletions.
diff --git a/tests/test__grib_idx.py b/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)

diff --git a/tests/test_xarray_backend.py b/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)