From 7ccb9c33eb2b72cf3ddb69a9bfda9ca624e0135a Mon Sep 17 00:00:00 2001
From: Alan Liddell <aliddell@chanzuckerberg.com>
Date: Fri, 17 Nov 2023 16:43:20 -0500
Subject: [PATCH] Add Zarr V3 tests. Add a note about environment variables in
 Zarr V3 tests.

---
 .github/workflows/release.yml |  4 +++
 .github/workflows/test_pr.yml |  2 ++
 README.md                     |  9 ++++++
 tests/test_basic.py           | 57 +++++++++++++++++++++++++++++++++++
 4 files changed, 72 insertions(+)

diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
index 8de8005..4e1e895 100644
--- a/.github/workflows/release.yml
+++ b/.github/workflows/release.yml
@@ -25,6 +25,10 @@ jobs:
     permissions:
       actions: write
 
+    env:
+      ZARR_V3_EXPERIMENTAL_API: 1
+      ZARR_V3_SHARDING: 1
+
     steps:
       - name: Cancel Previous Runs
         uses: styfle/cancel-workflow-action@0.11.0
diff --git a/.github/workflows/test_pr.yml b/.github/workflows/test_pr.yml
index ce2e764..d0a6e43 100644
--- a/.github/workflows/test_pr.yml
+++ b/.github/workflows/test_pr.yml
@@ -29,6 +29,8 @@ jobs:
       actions: write
     env:
       GH_TOKEN: ${{ github.token }}
+      ZARR_V3_EXPERIMENTAL_API: 1
+      ZARR_V3_SHARDING: 1
     steps:
       - name: Cancel Previous Runs
         uses: styfle/cancel-workflow-action@0.11.0
diff --git a/README.md b/README.md
index 154d106..d699d78 100644
--- a/README.md
+++ b/README.md
@@ -169,4 +169,13 @@ It depends on what you changed:
 - **acquire-video-runtime** (c/c++ code): `touch wrapper.h; maturin develop`
 - **rust code**: `maturin develop`
 
+### Zarr V3 tests are failing
+
+You should make sure that the following environment variables are set:
+
+```
+ZARR_V3_EXPERIMENTAL_API: 1
+ZARR_V3_SHARDING: 1
+```
+
 [napari]: https://github.com/napari/napari
diff --git a/tests/test_basic.py b/tests/test_basic.py
index d6b0b25..8ee688c 100644
--- a/tests/test_basic.py
+++ b/tests/test_basic.py
@@ -526,6 +526,63 @@ def test_write_zarr_multiscale(
         image = downscale_local_mean(image, (2, 2)).astype(np.uint8)
 
 
+@pytest.mark.parametrize(
+    ("number_of_frames", "expected_number_of_chunks", "codec"),
+    [
+        (64, 4, None),
+        (64, 4, "zstd"),
+        (65, 8, None),  # rollover
+        (65, 8, "lz4"),  # rollover
+    ],
+)
+def test_write_zarr_v3(
+    runtime: acquire.Runtime,
+    request: pytest.FixtureRequest,
+    number_of_frames: int,
+    expected_number_of_chunks: int,
+    codec: Optional[str],
+):
+    dm = runtime.device_manager()
+
+    p = runtime.get_configuration()
+    p.video[0].camera.identifier = dm.select(
+        DeviceKind.Camera, "simulated.*empty.*"
+    )
+
+    p.video[0].camera.settings.shape = (1920, 1080)
+    p.video[0].camera.settings.exposure_time_us = 1e4
+    p.video[0].camera.settings.pixel_type = acquire.SampleType.U8
+    p.video[0].storage.identifier = dm.select(
+        DeviceKind.Storage,
+        f"ZarrV3Blosc1{codec.capitalize()}ByteShuffle" if codec else "ZarrV3",
+    )
+    p.video[0].storage.settings.filename = f"{request.node.name}.zarr"
+    p.video[0].max_frame_count = number_of_frames
+
+    p.video[0].storage.settings.chunk_dims_px.width = 1920 // 2
+    p.video[0].storage.settings.chunk_dims_px.height = 1080 // 2
+    p.video[0].storage.settings.chunk_dims_px.planes = 64
+
+    runtime.set_configuration(p)
+
+    runtime.start()
+    runtime.stop()
+
+    store = zarr.DirectoryStoreV3(p.video[0].storage.settings.filename)
+    group = zarr.open(store=store, mode="r")
+    data = group["0"]
+
+    assert data.chunks == (64, 1, 1080 // 2, 1920 // 2)
+
+    assert data.shape == (
+        number_of_frames,
+        1,
+        p.video[0].camera.settings.shape[1],
+        p.video[0].camera.settings.shape[0],
+    )
+    assert data.nchunks == expected_number_of_chunks
+
+
 @pytest.mark.skip(
     reason="Runs into memory limitations on github ci."
     + " See https://github.com/acquire-project/cpx/issues/147"