parametrize test over async

Author: TomNicholas

xarray/tests/test_variable.py

@@ -2897,31 +2897,34 @@ def setUp(self):
         self.d = np.random.random((10, 3)).astype(np.float64)
         self.cat = PandasExtensionArray(pd.Categorical(["a", "b"] * 5))
 
-    async def check_orthogonal_indexing(self, v):
+    async def check_orthogonal_indexing(self, v, load_async):
         expected = self.d[[8, 3]][:, [2, 1]]
 
-        result = v.isel(x=[8, 3], y=[2, 1])
-        assert np.allclose(result, expected)
+        if load_async:
+            result = await v.isel(x=[8, 3], y=[2, 1]).load_async()
+        else:
+            result = v.isel(x=[8, 3], y=[2, 1])
 
-        result = await v.isel(x=[8, 3], y=[2, 1]).load_async()
         assert np.allclose(result, expected)
 
-    async def check_vectorized_indexing(self, v):
+    async def check_vectorized_indexing(self, v, load_async):
         ind_x = Variable("z", [0, 2])
         ind_y = Variable("z", [2, 1])
         expected = self.d[ind_x, ind_y]
 
-        result = v.isel(x=ind_x, y=ind_y).load()
-        assert np.allclose(result, expected)
+        if load_async:
+            result = await v.isel(x=ind_x, y=ind_y).load_async()
+        else:
+            result = v.isel(x=ind_x, y=ind_y).load()
 
-        result = await v.isel(x=ind_x, y=ind_y).load_async()
         assert np.allclose(result, expected)
 
     @pytest.mark.asyncio
-    async def test_NumpyIndexingAdapter(self):
+    @pytest.mark.parametrize("load_async", [True, False])
+    async def test_NumpyIndexingAdapter(self, load_async):
         v = Variable(dims=("x", "y"), data=NumpyIndexingAdapter(self.d))
-        await self.check_orthogonal_indexing(v)
-        await self.check_vectorized_indexing(v)
+        await self.check_orthogonal_indexing(v, load_async)
+        await self.check_vectorized_indexing(v, load_async)
         # could not doubly wrapping
         with pytest.raises(TypeError, match=r"NumpyIndexingAdapter only wraps "):
             v = Variable(
@@ -2937,57 +2940,61 @@ def test_extension_array_duck_indexed(self):
         assert (lazy[[0, 1, 5]] == ["a", "b", "b"]).all()
 
     @pytest.mark.asyncio
-    async def test_LazilyIndexedArray(self):
+    @pytest.mark.parametrize("load_async", [True, False])
+    async def test_LazilyIndexedArray(self, load_async):
         v = Variable(dims=("x", "y"), data=LazilyIndexedArray(self.d))
-        await self.check_orthogonal_indexing(v)
-        await self.check_vectorized_indexing(v)
+        await self.check_orthogonal_indexing(v, load_async)
+        await self.check_vectorized_indexing(v, load_async)
         # doubly wrapping
         v = Variable(
             dims=("x", "y"),
             data=LazilyIndexedArray(LazilyIndexedArray(self.d)),
         )
-        await self.check_orthogonal_indexing(v)
+        await self.check_orthogonal_indexing(v, load_async)
         # hierarchical wrapping
         v = Variable(
             dims=("x", "y"), data=LazilyIndexedArray(NumpyIndexingAdapter(self.d))
         )
-        await self.check_orthogonal_indexing(v)
+        await self.check_orthogonal_indexing(v, load_async)
 
     @pytest.mark.asyncio
-    async def test_CopyOnWriteArray(self):
+    @pytest.mark.parametrize("load_async", [True, False])
+    async def test_CopyOnWriteArray(self, load_async):
         v = Variable(dims=("x", "y"), data=CopyOnWriteArray(self.d))
-        await self.check_orthogonal_indexing(v)
-        await self.check_vectorized_indexing(v)
+        await self.check_orthogonal_indexing(v, load_async)
+        await self.check_vectorized_indexing(v, load_async)
         # doubly wrapping
         v = Variable(dims=("x", "y"), data=CopyOnWriteArray(LazilyIndexedArray(self.d)))
-        await self.check_orthogonal_indexing(v)
-        await self.check_vectorized_indexing(v)
+        await self.check_orthogonal_indexing(v, load_async)
+        await self.check_vectorized_indexing(v, load_async)
 
     @pytest.mark.asyncio
-    async def test_MemoryCachedArray(self):
+    @pytest.mark.parametrize("load_async", [True, False])
+    async def test_MemoryCachedArray(self, load_async):
         v = Variable(dims=("x", "y"), data=MemoryCachedArray(self.d))
-        await self.check_orthogonal_indexing(v)
-        await self.check_vectorized_indexing(v)
+        await self.check_orthogonal_indexing(v, load_async)
+        await self.check_vectorized_indexing(v, load_async)
         # doubly wrapping
         v = Variable(dims=("x", "y"), data=CopyOnWriteArray(MemoryCachedArray(self.d)))
-        await self.check_orthogonal_indexing(v)
-        await self.check_vectorized_indexing(v)
+        await self.check_orthogonal_indexing(v, load_async)
+        await self.check_vectorized_indexing(v, load_async)
 
     @requires_dask
     @pytest.mark.asyncio
-    async def test_DaskIndexingAdapter(self):
+    @pytest.mark.parametrize("load_async", [True, False])
+    async def test_DaskIndexingAdapter(self, load_async):
         import dask.array as da
 
         dask_array = da.asarray(self.d)
         v = Variable(dims=("x", "y"), data=DaskIndexingAdapter(dask_array))
-        await self.check_orthogonal_indexing(v)
-        await self.check_vectorized_indexing(v)
+        await self.check_orthogonal_indexing(v, load_async)
+        await self.check_vectorized_indexing(v, load_async)
         # doubly wrapping
         v = Variable(
             dims=("x", "y"), data=CopyOnWriteArray(DaskIndexingAdapter(dask_array))
         )
-        await self.check_orthogonal_indexing(v)
-        await self.check_vectorized_indexing(v)
+        await self.check_orthogonal_indexing(v, load_async)
+        await self.check_vectorized_indexing(v, load_async)
 
 
 def test_clip(var):