Use Cubed version of array assertion functions in tests (#802)

Author: tomwhite

cubed/tests/array/test_nan_functions.py

@@ -1,9 +1,9 @@
 import numpy as np
 import pytest
-from numpy.testing import assert_array_equal
 
 import cubed
 import cubed.array_api as xp
+from cubed._testing import assert_array_equal
 
 
 @pytest.fixture

cubed/tests/array/test_overlap.py

@@ -1,8 +1,8 @@
 import numpy as np
-from numpy.testing import assert_array_equal
 
 import cubed
 import cubed.array_api as xp
+from cubed._testing import assert_array_equal
 
 
 def test_map_overlap_1d():

cubed/tests/array/test_pad.py

@@ -1,9 +1,9 @@
 import numpy as np
 import pytest
-from numpy.testing import assert_array_equal
 
 import cubed
 import cubed.array_api as xp
+from cubed._testing import assert_array_equal
 
 
 @pytest.fixture

cubed/tests/primitive/test_blockwise.py

@@ -1,7 +1,7 @@
 import numpy as np
 import pytest
-from numpy.testing import assert_array_equal
 
+from cubed._testing import assert_array_equal
 from cubed.backend_array_api import namespace as nxp
 from cubed.primitive.blockwise import (
     blockwise,

cubed/tests/primitive/test_rechunk.py

@@ -1,8 +1,8 @@
 import numpy as np
 import pytest
 import zarr
-from numpy.testing import assert_array_equal
 
+from cubed._testing import assert_array_equal
 from cubed.primitive.rechunk import rechunk
 from cubed.runtime.executors.local import SingleThreadedExecutor
 from cubed.storage.backend import open_backend_array

cubed/tests/test_core.py

@@ -5,11 +5,11 @@
 import dill
 import numpy as np
 import pytest
-from numpy.testing import assert_array_equal
 
 import cubed
 import cubed.array_api as xp
 import cubed.random
+from cubed._testing import assert_array_equal
 from cubed.array_api.dtypes import _floating_dtypes
 from cubed.backend_array_api import namespace as nxp
 from cubed.core.ops import general_blockwise, merge_chunks, partial_reduce, tree_reduce

cubed/tests/test_executor_features.py

@@ -5,11 +5,11 @@
 import numpy as np
 import psutil
 import pytest
-from numpy.testing import assert_array_equal
 
 import cubed
 import cubed.array_api as xp
 import cubed.random
+from cubed._testing import assert_array_equal
 from cubed.diagnostics import ProgressBar
 from cubed.diagnostics.history import HistoryCallback
 from cubed.diagnostics.mem_usage import MemoryVisualizationCallback

cubed/tests/test_groupby.py

@@ -1,9 +1,9 @@
 import numpy as np
 import numpy_groupies as npg
 import pytest
-from numpy.testing import assert_array_equal
 
 import cubed.array_api as xp
+from cubed._testing import assert_array_equal
 from cubed.backend_array_api import namespace as nxp
 from cubed.core.groupby import (
     _get_chunks_for_groups,

cubed/tests/test_gufunc.py

@@ -1,10 +1,10 @@
 import numpy as np
 import pytest
-from numpy.testing import assert_allclose, assert_equal
 
 import cubed
 import cubed.array_api as xp
 from cubed import apply_gufunc
+from cubed._testing import assert_allclose, assert_array_equal
 from cubed.backend_array_api import namespace as nxp
 
 
@@ -35,7 +35,7 @@ def add(x, y):
     a = cubed.from_array(np.array([1, 2, 3]), chunks=2, spec=spec)
     b = cubed.from_array(np.array([1, 2, 3]), chunks=2, spec=spec)
     z = apply_gufunc(add, "(),()->()", a, b, output_dtypes=a.dtype)
-    assert_equal(z, np.array([2, 4, 6]))
+    assert_array_equal(z, np.array([2, 4, 6]))
 
 
 def test_apply_gufunc_elemwise_01_non_cubed_input(spec):
@@ -45,7 +45,7 @@ def add(x, y):
     a = cubed.from_array(np.array([1, 2, 3]), chunks=3, spec=spec)
     b = np.array([1, 2, 3])
     z = apply_gufunc(add, "(),()->()", a, b, output_dtypes=a.dtype)
-    assert_equal(z, np.array([2, 4, 6]))
+    assert_array_equal(z, np.array([2, 4, 6]))
 
 
 def test_apply_gufunc_elemwise_loop(spec):
@@ -56,7 +56,7 @@ def foo(x):
     a = cubed.from_array(np.array([1, 2, 3]), chunks=2, spec=spec)
     z = apply_gufunc(foo, "()->()", a, output_dtypes=int)
     assert z.chunks == ((2, 1),)
-    assert_equal(z, np.array([2, 4, 6]))
+    assert_array_equal(z, np.array([2, 4, 6]))
 
 
 def test_apply_gufunc_elemwise_core(spec):
@@ -67,7 +67,7 @@ def foo(x):
     a = cubed.from_array(np.array([1, 2, 3]), chunks=3, spec=spec)
     z = apply_gufunc(foo, "(i)->(i)", a, output_dtypes=int)
     assert z.chunks == ((3,),)
-    assert_equal(z, np.array([2, 4, 6]))
+    assert_array_equal(z, np.array([2, 4, 6]))
 
 
 def test_gufunc_two_inputs(spec):
@@ -77,7 +77,7 @@ def foo(x, y):
     a = xp.ones((2, 3), chunks=100, dtype=int, spec=spec)
     b = xp.ones((3, 4), chunks=100, dtype=int, spec=spec)
     x = apply_gufunc(foo, "(i,j),(j,k)->(i,k)", a, b, output_dtypes=int)
-    assert_equal(x, 3 * np.ones((2, 4), dtype=int))
+    assert_array_equal(x, 3 * np.ones((2, 4), dtype=int))
 
 
 def test_apply_gufunc_axes_two_kept_coredims(spec):
@@ -101,7 +101,7 @@ def foo(x):
     a = cubed.from_array(np.array([1, 2, 3, 4, 5], dtype=int), spec=spec)
     x = apply_gufunc(foo, "()->(i_0)", a, output_dtypes=int, output_sizes={"i_0": 3})
     assert x.chunks == ((5,), (3,))
-    assert_equal(
+    assert_array_equal(
         x,
         np.array(
             [

cubed/tests/test_icechunk.py

@@ -3,11 +3,11 @@
 import numpy as np
 import pytest
 import zarr
-from numpy.testing import assert_array_equal
 
 import cubed
 import cubed.array_api as xp
 import cubed.random
+from cubed._testing import assert_array_equal
 from cubed.tests.utils import MAIN_EXECUTORS
 
 icechunk = pytest.importorskip("icechunk")