From 41c43b8f3300c3d78feb532927961da7a434ac37 Mon Sep 17 00:00:00 2001
From: Igor Tatarnikov <igor.tatarnikov@gmail.com>
Date: Thu, 14 Mar 2024 11:02:11 +0000
Subject: [PATCH] Add tests for get_metadata_for_zarr

---
 brainglobe_stitch/image_mosaic.py    |  9 ++++--
 tests/test_unit/test_image_mosaic.py | 46 +++++++++++++++++++++++++---
 2 files changed, 47 insertions(+), 8 deletions(-)

diff --git a/brainglobe_stitch/image_mosaic.py b/brainglobe_stitch/image_mosaic.py
index dd622ea..38733c9 100644
--- a/brainglobe_stitch/image_mosaic.py
+++ b/brainglobe_stitch/image_mosaic.py
@@ -594,7 +594,10 @@ def fuse(
             self._fuse_to_bdv_h5(output_path, fused_image_shape)
 
     def _fuse_to_zarr(
-        self, output_path: Path, fused_image_shape: Tuple[int, ...]
+        self,
+        output_path: Path,
+        fused_image_shape: Tuple[int, ...],
+        pyramid_depth: int = 6,
     ) -> None:
         """
         Fuse the tiles in the ImageMosaic into a single image and save it as a
@@ -613,7 +616,7 @@ def _fuse_to_zarr(
         chunk_shape: Tuple[int, ...] = (256, 256, 256)
 
         transformation_metadata, axes_metadata = self.get_metadata_for_zarr(
-            pyramid_depth=6
+            pyramid_depth=pyramid_depth
         )
 
         if self.num_channels > 1:
@@ -846,7 +849,7 @@ def _fuse_to_bdv_h5(
         output_file.close()
 
     def get_metadata_for_zarr(
-        self, pyramid_depth: int = 5
+        self, pyramid_depth: int = 6
     ) -> Tuple[List[List[Dict]], List[Dict]]:
         """
         Prepare the metadata for a zarr file.
diff --git a/tests/test_unit/test_image_mosaic.py b/tests/test_unit/test_image_mosaic.py
index 9405e8a..1cd3973 100644
--- a/tests/test_unit/test_image_mosaic.py
+++ b/tests/test_unit/test_image_mosaic.py
@@ -349,13 +349,49 @@ def test_fuse(
         mock_interpolate_overlaps.assert_not_called()
 
 
-def test_fuse_to_zarr():
-    pass
+@pytest.mark.parametrize(
+    "pyramid_depth, num_channels",
+    [(1, 1), (1, 2), (1, 5), (2, 1), (2, 2), (2, 5), (3, 1), (3, 2), (3, 5)],
+)
+def test_get_metadata_for_zarr(image_mosaic, pyramid_depth, num_channels):
+    temp_num_channels = image_mosaic.num_channels
+    image_mosaic.num_channels = num_channels
 
+    metadata, axes = image_mosaic.get_metadata_for_zarr(pyramid_depth)
 
-def test_fuse_to_bdv_h5():
-    pass
+    if num_channels > 1:
+        assert len(axes) == 4
+        assert axes[0]["name"] == "c"
+        assert axes[0]["type"] == "channel"
+        axes.pop(0)
+    else:
+        assert len(axes) == 3
+
+    expected_axes_names = ["z", "y", "x"]
 
+    for idx, axes in enumerate(axes):
+        assert axes["name"] == expected_axes_names[idx]
+        assert axes["type"] == "space"
+        assert axes["unit"] == "micrometer"
 
-def test_get_metadata_for_zarr():
+    for idx, transformation in enumerate(metadata):
+        assert transformation[0]["type"] == "scale"
+        expected_scale = [
+            image_mosaic.z_resolution,
+            image_mosaic.x_y_resolution * 2**idx,
+            image_mosaic.x_y_resolution * 2**idx,
+        ]
+        if num_channels > 1:
+            expected_scale.insert(0, 1)
+
+        assert transformation[0]["scale"] == expected_scale
+
+    image_mosaic.num_channels = temp_num_channels
+
+
+def test_fuse_to_zarr(image_mosaic):
+    image_mosaic.xml_path.parent / "fused.zarr"
+
+
+def test_fuse_to_bdv_h5():
     pass