Skip to content

Commit f67a68e

Browse files
committed
Adding test checking that existing host memory codecs use the gpu_buffer_prototype appropriately
1 parent ccd266e commit f67a68e

File tree

1 file changed

+31
-0
lines changed

1 file changed

+31
-0
lines changed

tests/v3/test_buffer.py

+31
Original file line numberDiff line numberDiff line change
@@ -170,3 +170,34 @@ async def test_codecs_use_of_prototype():
170170
got = await a.getitem(selection=(slice(0, 10), slice(0, 10)), prototype=my_prototype)
171171
assert isinstance(got, MyNDArrayLike)
172172
assert np.array_equal(expect, got)
173+
174+
175+
@pytest.mark.skipif(cp is None, reason="requires cupy")
176+
@pytest.mark.asyncio
177+
async def test_codecs_use_of_gpu_prototype():
178+
expect = cp.zeros((10, 10), dtype="uint16", order="F")
179+
a = await AsyncArray.create(
180+
StorePath(MemoryStore(mode="w")) / "test_codecs_use_of_gpu_prototype",
181+
shape=expect.shape,
182+
chunk_shape=(5, 5),
183+
dtype=expect.dtype,
184+
fill_value=0,
185+
codecs=[
186+
TransposeCodec(order=(1, 0)),
187+
BytesCodec(),
188+
BloscCodec(),
189+
Crc32cCodec(),
190+
GzipCodec(),
191+
ZstdCodec(),
192+
],
193+
)
194+
expect[:] = cp.arange(100).reshape(10, 10)
195+
196+
await a.setitem(
197+
selection=(slice(0, 10), slice(0, 10)),
198+
value=expect[:],
199+
prototype=gpu_buffer_prototype,
200+
)
201+
got = await a.getitem(selection=(slice(0, 10), slice(0, 10)), prototype=gpu_buffer_prototype)
202+
assert isinstance(got, cp.ndarray)
203+
assert cp.array_equal(expect, got)

0 commit comments

Comments
 (0)