From f03283fe054c37293659c64773f2ee13d045c379 Mon Sep 17 00:00:00 2001
From: Erik O Gabrielsson <erik.o.gabrielsson@sectra.com>
Date: Wed, 30 Oct 2024 17:35:43 +0100
Subject: [PATCH] Test for ImageCoordinateSystemDicomSchema load

---
 .../dicom_schema/test_dicom_schema.py         | 43 ++++++++++++++++++-
 wsidicom/metadata/schema/dicom/image.py       |  6 +--
 2 files changed, 43 insertions(+), 6 deletions(-)

diff --git a/tests/metadata/dicom_schema/test_dicom_schema.py b/tests/metadata/dicom_schema/test_dicom_schema.py
index 35f8121..caa72f2 100644
--- a/tests/metadata/dicom_schema/test_dicom_schema.py
+++ b/tests/metadata/dicom_schema/test_dicom_schema.py
@@ -13,6 +13,7 @@
 #    limitations under the License.
 
 from datetime import datetime
+from typing import Optional
 
 import pytest
 from pydicom import Dataset
@@ -36,7 +37,7 @@
     LenseCode,
     LightPathFilterCode,
 )
-from wsidicom.geometry import PointMm, SizeMm
+from wsidicom.geometry import Orientation, PointMm, SizeMm
 from wsidicom.instance import ImageType
 from wsidicom.metadata import (
     Equipment,
@@ -65,7 +66,10 @@
 )
 from wsidicom.metadata.schema.dicom.defaults import Defaults
 from wsidicom.metadata.schema.dicom.equipment import EquipmentDicomSchema
-from wsidicom.metadata.schema.dicom.image import ImageDicomSchema
+from wsidicom.metadata.schema.dicom.image import (
+    ImageCoordinateSystemDicomSchema,
+    ImageDicomSchema,
+)
 from wsidicom.metadata.schema.dicom.label import LabelDicomSchema
 from wsidicom.metadata.schema.dicom.optical_path import (
     OpticalPathDicomSchema,
@@ -796,3 +800,38 @@ def test_deserialize_wsi_metadata_from_empty_dataset(
 
         # Assert
         assert isinstance(deserialized, WsiMetadata)
+
+    @pytest.mark.parametrize("origin", [PointMm(20.0, 30.0), None])
+    @pytest.mark.parametrize("orientation", [Orientation.from_rotation(90), None])
+    @pytest.mark.parametrize("z_offset", [1.0, None])
+    def test_deserialize_image_coordinate_system(
+        self,
+        origin: Optional[PointMm],
+        orientation: Optional[Orientation],
+        z_offset: Optional[float],
+    ):
+        # Arrange
+        dataset = Dataset()
+        origin_dataset = Dataset()
+        if origin is not None:
+            origin_dataset.XOffsetInSlideCoordinateSystem = origin.x
+            origin_dataset.YOffsetInSlideCoordinateSystem = origin.y
+        if z_offset is not None:
+            origin_dataset.ZOffsetInSlideCoordinateSystem = z_offset
+        if len(origin_dataset) > 0:
+            dataset.TotalPixelMatrixOriginSequence = [origin_dataset]
+        if orientation is not None:
+            dataset.ImageOrientationSlide = list(orientation.values)
+
+        schema = ImageCoordinateSystemDicomSchema()
+        # Act
+        deserialized = schema.load(dataset)
+
+        # Assert
+        if origin is not None and orientation is not None:
+            assert isinstance(deserialized, ImageCoordinateSystem)
+            assert deserialized.origin == origin
+            assert deserialized.rotation == orientation.rotation
+            assert deserialized.z_offset == z_offset
+        else:
+            assert deserialized is None
diff --git a/wsidicom/metadata/schema/dicom/image.py b/wsidicom/metadata/schema/dicom/image.py
index 5aa527d..954cca9 100644
--- a/wsidicom/metadata/schema/dicom/image.py
+++ b/wsidicom/metadata/schema/dicom/image.py
@@ -83,7 +83,7 @@ def load_type(self) -> Type[ImageCoordinateSystem]:
     def load(self, dataset: Dataset, **kwargs) -> Optional[ImageCoordinateSystem]:
         try:
             return super().load(dataset, **kwargs)
-        except (TypeError, AttributeError):
+        except (TypeError, AttributeError, KeyError, IndexError):
             return None
 
     @post_load
@@ -91,9 +91,7 @@ def post_load(
         self, data: Dict[str, Any], **kwargs
     ) -> Optional[ImageCoordinateSystem]:
         """Post load hook to handle separation of xy and z offset."""
-        origin: Optional[Tuple[ImageCoordinateSystem, float]] = data.pop("origin", None)
-        if origin is None:
-            return None
+        origin: Tuple[ImageCoordinateSystem, float] = data.pop("origin")
         return super().post_load(
             {"origin": (origin[0]), "rotation": data["rotation"], "z_offset": origin[1]}
         )