Add test for apply_ufunc

Author: tomwhite

cubed_xarray/tests/test_wrapping.py

@@ -4,6 +4,7 @@
 import pytest
 import xarray as xr
 from cubed.runtime.create import create_executor
+from numpy.testing import assert_array_equal
 from xarray.namedarray.parallelcompat import list_chunkmanagers
 from xarray.tests import assert_allclose, create_test_data
 
@@ -24,6 +25,17 @@ def executor(request):
     return request.param
 
 
+def assert_identical(a, b):
+    """A version of this function which accepts numpy arrays"""
+    __tracebackhide__ = True
+    from xarray.testing import assert_identical as assert_identical_
+
+    if hasattr(a, "identical"):
+        assert_identical_(a, b)
+    else:
+        assert_array_equal(a, b)
+
+
 class TestDiscoverCubedManager:
     def test_list_cubedmanager(self):
         chunkmanagers = list_chunkmanagers()
@@ -72,3 +84,31 @@ def test_dataset_accessor_visualize(tmp_path):
     assert not (tmp_path / "cubed.svg").exists()
     ds.cubed.visualize(filename=tmp_path / "cubed")
     assert (tmp_path / "cubed.svg").exists()
+
+
+def identity(x):
+    return x
+
+
+# based on test_apply_dask_parallelized_one_arg
+def test_apply_ufunc_parallelized_one_arg():
+    array = cubed.ones((2, 2), chunks=(1, 1))
+    data_array = xr.DataArray(array, dims=("x", "y"))
+
+    def parallel_identity(x):
+        return xr.apply_ufunc(
+            identity,
+            x,
+            output_dtypes=[x.dtype],
+            dask="parallelized",
+            dask_gufunc_kwargs={"allow_rechunk": False},
+        )
+
+    actual = parallel_identity(data_array)
+    assert isinstance(actual.data, cubed.Array)
+    assert actual.data.chunks == array.chunks
+    assert_identical(data_array, actual)
+
+    computed = data_array.compute()
+    actual = parallel_identity(computed)
+    assert_identical(computed, actual)