From b58d55035483e1d6c9cd552ab293d1ce7d26a26b Mon Sep 17 00:00:00 2001 From: snibbor Date: Sun, 25 Aug 2024 21:39:21 -0400 Subject: [PATCH] test for correct chunk size after writing to ZipStore --- tests/test_writer.py | 54 ++++++++++++++++++++++++++++++++++++-------- 1 file changed, 45 insertions(+), 9 deletions(-) diff --git a/tests/test_writer.py b/tests/test_writer.py index 03beba65..b1419350 100644 --- a/tests/test_writer.py +++ b/tests/test_writer.py @@ -157,9 +157,27 @@ def test_write_image_zipstore(self, array_constructor, storage_options): shape = (3, 300, 300) data = self.create_data(shape) - data = array_constructor(data) + if array_constructor == da.from_array: + data = array_constructor(data, chunks=(1, 50, 50)) + else: + data = array_constructor(data) write_image(data, group, axes="cyx", storage_options=storage_options) + self.zip_store.close() + # load the data from the ZipStore + store = zarr.storage.ZipStore(self.zip_path, mode="r") + group = zarr.open_group(store=store) + # check the chunksize of the array + if array_constructor == da.from_array: + if storage_options is None: + assert group[0].chunks == (1, 50, 50) + else: + assert group[0].chunks == storage_options["chunks"] + else: + if storage_options is None: + pass + else: + assert group[0].chunks == storage_options["chunks"] @pytest.mark.parametrize("array_constructor", [np.array, da.from_array]) @pytest.mark.parametrize("storage_options", [None, {"chunks": (1, 100, 100)}]) @@ -168,14 +186,32 @@ def test_write_multiscale_zipstore(self, array_constructor, storage_options): group = zarr.group(store=self.zip_store, overwrite=True) shape = (3, 300, 300) - data1 = self.create_data(shape) - data1 = array_constructor(data1) - data2 = self.create_data(shape) - data2 = array_constructor(data2) - - write_multiscale( - [data1, data2], group, axes="cyx", storage_options=storage_options - ) + data_arrs = [] + for i in range(2): + data = self.create_data(shape) + if array_constructor == da.from_array: + data = array_constructor(data, chunks=(1, 50, 50)) + else: + data = array_constructor(data) + data_arrs.append(data.copy()) + + write_multiscale(data_arrs, group, axes="cyx", storage_options=storage_options) + self.zip_store.close() + # load the data from the ZipStore + store = zarr.storage.ZipStore(self.zip_path, mode="r") + group = zarr.open_group(store=store) + # check the chunksize of the array + for i in range(2): + if array_constructor == da.from_array: + if storage_options is None: + assert group[i].chunks == (1, 50, 50) + else: + assert group[i].chunks == storage_options["chunks"] + else: + if storage_options is None: + pass + else: + assert group[i].chunks == storage_options["chunks"] @pytest.mark.parametrize("read_from_zarr", [True, False]) @pytest.mark.parametrize("compute", [True, False])