From 72f684eff6de06b31b3fba00c2d1d8e283dc93cc Mon Sep 17 00:00:00 2001
From: emfdavid <stu3b3+emfdavid@gmail.com>
Date: Mon, 16 Dec 2024 14:40:50 -0500
Subject: [PATCH] Fix tests and black format

---
 tests/test__grib_idx.py      | 27 +++++++++++++++++++++++----
 tests/test_xarray_backend.py |  1 +
 2 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/tests/test__grib_idx.py b/tests/test__grib_idx.py
index 7d34ecc9..20fde1f3 100644
--- 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
index 8846ea9d..a2944a02 100644
--- 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)