test fixes

NREL · Jan 8, 2025 · fd9055c · fd9055c
1 parent 0e246a2
commit fd9055c
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/sup3r/utilities/interpolation.py b/sup3r/utilities/interpolation.py
@@ -48,7 +48,7 @@ def get_level_masks(cls, lev_array, level):
             da.arange(lev_array.shape[-1]), lev_array.shape
         )
         mask1 = lev_indices == argmin1
-        lev_diff = da.where(mask1, np.inf, lev_diff)
+        lev_diff = da.abs(da.ma.masked_array(lev_array, mask1) - level)
         argmin2 = da.argmin(lev_diff, axis=-1, keepdims=True)
         mask2 = lev_indices == argmin2
         return mask1, mask2

diff --git a/tests/data_handlers/test_dh_nc_cc.py b/tests/data_handlers/test_dh_nc_cc.py
@@ -69,6 +69,7 @@ def test_reload_cache():
             target=target,
             shape=(20, 20),
             cache_kwargs=cache_kwargs,
+            interp_kwargs={'include_single_levels': True}
         )
 
         # reload from cache
@@ -80,7 +81,9 @@ def test_reload_cache():
             cache_kwargs=cache_kwargs,
         )
         assert all(f in cached for f in features)
-        assert np.array_equal(handler.as_array(), cached.as_array())
+        harr = handler.as_array().compute()
+        carr = cached.as_array().compute()
+        assert np.array_equal(harr, carr)
 
 
 @pytest.mark.parametrize(