handle warnings in gpu tests

Author: d-v-b

tests/test_buffer.py

@@ -13,6 +13,7 @@
 from zarr.codecs.gzip import GzipCodec
 from zarr.codecs.transpose import TransposeCodec
 from zarr.codecs.zstd import ZstdCodec
+from zarr.errors import ZarrUserWarning
 from zarr.storage import MemoryStore, StorePath
 from zarr.testing.buffer import (
     NDBufferUsingTestNDArrayLike,
@@ -138,12 +139,13 @@ async def test_codecs_use_of_gpu_prototype() -> None:
         filters=[TransposeCodec(order=(1, 0))],
     )
     expect[:] = cp.arange(100).reshape(10, 10)
-
-    await a.setitem(
-        selection=(slice(0, 10), slice(0, 10)),
-        value=expect[:],
-        prototype=gpu.buffer_prototype,
-    )
+    msg = "Creating a zarr.buffer.gpu.Buffer with an array that does not support the __cuda_array_interface__ for zero-copy transfers, falling back to slow copy based path"
+    with pytest.warns(ZarrUserWarning, match=msg):
+        await a.setitem(
+            selection=(slice(0, 10), slice(0, 10)),
+            value=expect[:],
+            prototype=gpu.buffer_prototype,
+        )
     got = await a.getitem(selection=(slice(0, 10), slice(0, 10)), prototype=gpu.buffer_prototype)
     assert isinstance(got, cp.ndarray)
     assert cp.array_equal(expect, got)
@@ -164,12 +166,13 @@ async def test_sharding_use_of_gpu_prototype() -> None:
             fill_value=0,
         )
         expect[:] = cp.arange(100).reshape(10, 10)
-
-        await a.setitem(
-            selection=(slice(0, 10), slice(0, 10)),
-            value=expect[:],
-            prototype=gpu.buffer_prototype,
-        )
+        msg = "Creating a zarr.buffer.gpu.Buffer with an array that does not support the __cuda_array_interface__ for zero-copy transfers, falling back to slow copy based path"
+        with pytest.warns(ZarrUserWarning, match=msg):
+            await a.setitem(
+                selection=(slice(0, 10), slice(0, 10)),
+                value=expect[:],
+                prototype=gpu.buffer_prototype,
+            )
         got = await a.getitem(
             selection=(slice(0, 10), slice(0, 10)), prototype=gpu.buffer_prototype
         )

tests/test_store/test_memory.py

@@ -9,6 +9,7 @@
 
 import zarr
 from zarr.core.buffer import Buffer, cpu, gpu
+from zarr.errors import ZarrUserWarning
 from zarr.storage import GpuMemoryStore, MemoryStore
 from zarr.testing.store import StoreTests
 from zarr.testing.utils import gpu_test
@@ -130,6 +131,8 @@ def test_from_dict(self) -> None:
             "a": gpu.Buffer.from_bytes(b"aaaa"),
             "b": cpu.Buffer.from_bytes(b"bbbb"),
         }
-        result = GpuMemoryStore.from_dict(d)
+        msg = "Creating a zarr.buffer.gpu.Buffer with an array that does not support the __cuda_array_interface__ for zero-copy transfers, falling back to slow copy based path"
+        with pytest.warns(ZarrUserWarning, match=msg):
+            result = GpuMemoryStore.from_dict(d)
         for v in result._store_dict.values():
             assert type(v) is gpu.Buffer