Cubed xarray tests

ci/requirements/environment.yaml

@@ -10,3 +10,5 @@ dependencies:
   - hypothesis
   - xarray
   - numpy
+  - cubed
+  - cubed-xarray

pyproject.toml

@@ -28,7 +28,7 @@ exclude = [
     "__pycache__",
     "docs",
 ]
-target-version = "py312"
+target-version = "py311"
 
 extend-include = ["*.ipynb"]
 line-length = 100

xarray_array_testing/creation.py

@@ -1,11 +1,12 @@
 import hypothesis.strategies as st
 import xarray.testing.strategies as xrst
-from hypothesis import given
+from hypothesis import HealthCheck, given, settings
 
 from xarray_array_testing.base import DuckArrayTestMixin
 
 
 class CreationTests(DuckArrayTestMixin):
+    @settings(suppress_health_check=[HealthCheck.differing_executors])
     @given(st.data())
     def test_create_variable(self, data):
         variable = data.draw(xrst.variables(array_strategy_fn=self.array_strategy_fn))

xarray_array_testing/reduction.py

@@ -3,7 +3,7 @@
 import hypothesis.strategies as st
 import pytest
 import xarray.testing.strategies as xrst
-from hypothesis import given
+from hypothesis import HealthCheck, given, note, settings
 
 from xarray_array_testing.base import DuckArrayTestMixin
 
@@ -13,15 +13,22 @@ class ReductionTests(DuckArrayTestMixin):
     def expected_errors(op, **parameters):
         return nullcontext()
 
+    # TODO understand the differing executors health check error
+    @settings(suppress_health_check=[HealthCheck.differing_executors])
     @pytest.mark.parametrize("op", ["mean", "sum", "prod", "std", "var"])
     @given(st.data())
     def test_variable_numerical_reduce(self, op, data):
         variable = data.draw(xrst.variables(array_strategy_fn=self.array_strategy_fn))
 
+        note(f"note: {variable}")
+
         with self.expected_errors(op, variable=variable):
             # compute using xr.Variable.<OP>()
             actual = getattr(variable, op)().data
             # compute using xp.<OP>(array)
             expected = getattr(self.xp, op)(variable.data)
 
-        self.assert_equal(actual, expected)
+            assert isinstance(
+                actual, self.array_type
+            ), f"expected {self.array_type} but got {type(actual)}"
+            self.assert_equal(actual, expected)

xarray_array_testing/tests/test_cubed.py

@@ -0,0 +1,89 @@
+from contextlib import AbstractContextManager as ContextManager
+from contextlib import nullcontext
+
+import cubed
+import cubed.random
+import hypothesis.strategies as st
+import numpy as np
+import numpy.testing as npt
+import pytest
+from hypothesis import note
+
+from xarray_array_testing.base import DuckArrayTestMixin
+from xarray_array_testing.creation import CreationTests
+from xarray_array_testing.reduction import ReductionTests
+
+
+def cubed_random_array(shape: tuple[int], dtype: np.dtype) -> cubed.Array:
+    """
+    Generates a random cubed array
+
+    Supports integer and float dtypes.
+    """
+    # TODO hypothesis doesn't like us using random inside strategies
+    rng = np.random.default_rng()
+
+    if np.issubdtype(dtype, np.integer):
+        arr = rng.integers(low=0, high=+3, size=shape, dtype=dtype)
+        return cubed.from_array(arr)
+    else:
+        # TODO generate general chunking pattern
+        ca = cubed.random.random(size=shape, chunks=shape)
+        return cubed.array_api.astype(ca, dtype)
+
+
+def random_cubed_arrays_fn(
+    *,
+    shape: tuple[int, ...],
+    dtype: np.dtype,
+) -> st.SearchStrategy[cubed.Array]:
+    return st.builds(cubed_random_array, shape=st.just(shape), dtype=st.just(dtype))
+
+
+class CubedTestMixin(DuckArrayTestMixin):
+    @property
+    def xp(self) -> type[cubed.array_api]:
+        return cubed.array_api
+
+    @property
+    def array_type(self) -> type[cubed.Array]:
+        return cubed.Array
+
+    @staticmethod
+    def array_strategy_fn(*, shape, dtype) -> st.SearchStrategy[cubed.Array]:
+        return random_cubed_arrays_fn(shape=shape, dtype=dtype)
+
+    @staticmethod
+    def assert_equal(a: cubed.Array, b: cubed.Array):
+        npt.assert_equal(a.compute(), b.compute())
+
+
+class TestCreationCubed(CreationTests, CubedTestMixin):
+    pass
+
+
+class TestReductionCubed(ReductionTests, CubedTestMixin):
+    @staticmethod
+    def expected_errors(op, **parameters) -> ContextManager:
+        var = parameters.get("variable")
+
+        xp = cubed.array_api
+
+        note(f"op = {op}")
+        note(f"dtype = {var.dtype}")
+        note(f"is_integer = {cubed.array_api.isdtype(var.dtype, 'integral')}")
+
+        if op == "mean" and xp.isdtype(
+            var.dtype, ("integral", "complex floating", np.dtype("float16"))
+        ):
+            return pytest.raises(
+                TypeError, match="Only real floating-point dtypes are allowed in mean"
+            )
+        elif xp.isdtype(var.dtype, np.dtype("float16")):
+            return pytest.raises(
+                TypeError, match="Only numeric dtypes are allowed in isnan"
+            )
+        elif op in {"var", "std"}:
+            pytest.skip(reason=f"cubed does not implement {op} yet")
+        else:
+            return nullcontext()