From 8e65cc8ed26fc350324d3008b161b8558367fdf1 Mon Sep 17 00:00:00 2001
From: Anish <anish.shrestha@nih.gov>
Date: Fri, 1 Sep 2023 09:33:17 -0400
Subject: [PATCH] Add tests

---
 requirements-dev.txt          |  2 +-
 test/__init__.py              |  0
 test/test_HedwigZarrImages.py | 34 ++++++++++++++++++++++++----------
 test/test_zarr_rechunk.py     | 31 +++++++++++++++++++++++++++++--
 4 files changed, 54 insertions(+), 13 deletions(-)
 create mode 100644 test/__init__.py

diff --git a/requirements-dev.txt b/requirements-dev.txt
index ae45a57..78591c2 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -5,4 +5,4 @@ flake8
 black
 wheel
 packaging
--rrequirements.txt
\ No newline at end of file
+-r requirements.txt
diff --git a/test/__init__.py b/test/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/test/test_HedwigZarrImages.py b/test/test_HedwigZarrImages.py
index 8b7acc5..6ceae94 100644
--- a/test/test_HedwigZarrImages.py
+++ b/test/test_HedwigZarrImages.py
@@ -35,10 +35,24 @@ def temp_zarr_path(request, tmp_path_factory, data_path):
 
 
 @pytest.mark.parametrize(
-    "zarr_name",
-    ["2013-1220-dA30_5-BSC-1_22_full_rec.zarr", "OM_P1_S1_ScanOnly_1k.zarr", "OM_P1_S1_ScanOnly_1k_tiff.zarr"],
+    "zarr_name, image_ext, array_shape, dims, shader_type, ngff_dims, shader_params",
+    [
+        (
+            "2013-1220-dA30_5-BSC-1_22_full_rec.zarr",
+            "mrc",
+            (1, 1, 512, 87, 512),
+            "XYZ",
+            "Grayscale",
+            "TCZYX",
+            {"window": [263, 413], "range": [-2664, 1677]},
+        ),
+        ("OM_P1_S1_ScanOnly_1k.zarr", "png", (1, 3, 1, 1024, 521), "XYC", "RGB", "TCZYX", {}),
+        ("OM_P1_S1_ScanOnly_1k_tiff.zarr", "tiff", (1, 3, 1, 1024, 521), "XYC", "RGB", "TCZYX", {}),
+    ],
 )
-def test_HedwigZarrImage_info(data_path, zarr_name):
+def test_HedwigZarrImage_info(
+    data_path, zarr_name, image_ext, array_shape, dims, shader_type, ngff_dims, shader_params
+):
     zi = HedwigZarrImages(data_path / zarr_name)
     keys = list(zi.get_series_keys())
     print(f"zarr groups: {keys}")
@@ -46,13 +60,13 @@ def test_HedwigZarrImage_info(data_path, zarr_name):
     print(zi.ome_xml_path)
 
     for k, z_img in zi.series():
-        print(f'image name: "{k}"')
-        print(f"\tarray shape: {z_img.shape}")
-        print(f"\tzarr path: {z_img.path}")
-        print(f"\tdims: {z_img.dims}")
-        print(f"\tshader type: {z_img.shader_type}")
-        print(f"\tNGFF dims: {z_img._ome_ngff_multiscale_dims()}")
-        print(f"\tshader params: {json.dumps(z_img.neuroglancer_shader_parameters())}")
+        assert image_ext in k
+        assert array_shape == z_img.shape
+        print(f"{z_img.path=}")
+        assert dims == z_img.dims
+        assert shader_type == z_img.shader_type
+        assert ngff_dims == z_img._ome_ngff_multiscale_dims()
+        assert shader_params == z_img.neuroglancer_shader_parameters()
 
 
 @pytest.mark.parametrize(
diff --git a/test/test_zarr_rechunk.py b/test/test_zarr_rechunk.py
index 08f0ef5..becde85 100644
--- a/test/test_zarr_rechunk.py
+++ b/test/test_zarr_rechunk.py
@@ -1,12 +1,39 @@
 from click.testing import CliRunner
-import pytools.zarr_rechunk
+import shutil
+from pathlib import Path
 import pytest
+import pytools.zarr_rechunk
+from pytools.HedwigZarrImages import HedwigZarrImages
 
 args = ["--help", "--version"]
 
 
 @pytest.mark.parametrize("cli_args", args)
-def test_mrc2nifti_main_help(cli_args):
+def test_zarr_rechunk_main_help(cli_args):
     runner = CliRunner()
     result = runner.invoke(pytools.zarr_rechunk.main, cli_args.split())
     assert not result.exception
+
+
+@pytest.mark.parametrize("chunk_size", [512, 64])
+def test_zarr_rechunk_main(image_ome_ngff, chunk_size):
+    runner = CliRunner()
+    with runner.isolated_filesystem():
+        local_path = image_ome_ngff.name
+        shutil.copytree(image_ome_ngff, local_path)
+
+        result = runner.invoke(pytools.zarr_rechunk.main, args=[str(local_path), "--chunk-size", chunk_size])
+        assert not result.exception
+
+        zi = HedwigZarrImages(Path(local_path).absolute(), read_only=False)
+        for k, hzi in zi.series():
+            for level in range(len(hzi._ome_ngff_multiscales()["datasets"])):
+                arr = hzi._ome_ngff_multiscale_get_array(level)
+                # check for expected chunking
+                for s, c, d in zip(arr.shape, arr.chunks, hzi._ome_ngff_multiscale_dims()):
+                    if d == "T":
+                        assert s == 1
+                    elif d == "C" or s < target_chunk:
+                        assert s == c
+                    else:
+                        assert c == target_chunk