test(data): add unit tests the data factory

Author: MattTheCuber

edge_mender/data_factory.py

@@ -1,66 +1,17 @@
+"""Contains a class with a set of test case data sets."""
+
 import cc3d
 import fill_voids
 import numpy as np
 from numpy.typing import NDArray
 
 
 class DataFactory:
-    @staticmethod
-    def kill_you() -> NDArray:
-        # Create initial data
-        data = np.zeros((10, 10, 10))
-        data[3:-3, 3:-3, 3:-3] = 1
-
-        # Case 1 - simple extrusion
-        data[2, 4, 4] = 1
-        data[2, 5, 5] = 1
-
-        # Case 2 - double extrusion
-        data[4, 4, 2] = 1
-        data[4, 4, 1] = 1
-        data[5, 5, 2] = 1
-        data[5, 5, 1] = 1
-
-        # Case 3 - triple extrusion
-        data[4, 7, 4] = 1
-        data[4, 8, 4] = 1
-        data[4, 9, 4] = 1
-        data[5, 7, 5] = 1
-        data[5, 8, 5] = 1
-        data[5, 9, 5] = 1
-
-        # Case 4 - stairs
-        data[4, 4, 7] = 1
-        data[4, 4, 8] = 1
-        data[4, 5, 7] = 1
-        data[5, 5, 7] = 1
-        data[5, 5, 8] = 1
-        data[5, 5, 9] = 1
-
-        # Case 5 - ceiling
-        data[4, 1, 4] = 1
-        data[4, 1, 5] = 1
-        data[4, 2, 4] = 1
-        data[5, 1, 4] = 1
-        data[5, 1, 5] = 1
-        data[5, 2, 5] = 1
-
-        # Case 6 - checkboard
-        data[7, 5, 4] = 1
-        data[7, 4, 5] = 1
-        data[8, 4, 4] = 1
-        data[8, 3, 4] = 1
-        data[8, 3, 3] = 1
-        data[7, 3, 3] = 1
-        data[8, 5, 5] = 1
-        data[8, 6, 5] = 1
-        data[8, 6, 6] = 1
-        data[7, 6, 6] = 1
-
-        return data
+    """A set of test cases for testing, evaluation, and demonstration."""
 
     @staticmethod
     def simple_extrusion() -> NDArray:
+        """Create a test case with a simple extrusion."""
         # Create initial data
         data = np.zeros((4, 4, 4))
 
@@ -75,6 +26,7 @@ def simple_extrusion() -> NDArray:
 
     @staticmethod
     def double_extrusion() -> NDArray:
+        """Create a test case with a double extrusion."""
         # Create initial data
         data = np.zeros((4, 5, 4))
 
@@ -89,6 +41,7 @@ def double_extrusion() -> NDArray:
 
     @staticmethod
     def triple_extrusion() -> NDArray:
+        """Create a test case with a triple extrusion."""
         # Create initial data
         data = np.zeros((4, 6, 4))
 
@@ -103,6 +56,7 @@ def triple_extrusion() -> NDArray:
 
     @staticmethod
     def stairs() -> NDArray:
+        """Create a test case with stairs."""
         # Create initial data
         data = np.zeros((4, 6, 4))
 
@@ -122,6 +76,7 @@ def stairs() -> NDArray:
 
     @staticmethod
     def ceiling() -> NDArray:
+        """Create a test case with a ceiling."""
         # Create initial data
         data = np.zeros((4, 5, 4))
 
@@ -139,6 +94,7 @@ def ceiling() -> NDArray:
 
     @staticmethod
     def double_tower_ceiling() -> NDArray:
+        """Create a test case with a double tower ceiling."""
         # Create initial data
         data = np.zeros((4, 5, 5))
 
@@ -165,6 +121,7 @@ def double_tower_ceiling() -> NDArray:
 
     @staticmethod
     def hanging_points() -> NDArray:
+        """Create a test case with hanging points."""
         # Create initial data
         data = np.zeros((5, 5, 3))
 
@@ -180,6 +137,7 @@ def hanging_points() -> NDArray:
 
     @staticmethod
     def checkerboard() -> NDArray:
+        """Create a test case with a checkerboard pattern."""
         # Create initial data
         data = np.zeros((4, 5, 5))
 
@@ -206,7 +164,7 @@ def checkerboard() -> NDArray:
 
     @staticmethod
     def hole() -> NDArray:
-        """faces_to_flip=[44, 45, 54, 55, 66, 67, 68, 69, 70, 71]"""
+        """Create a test case with a hole in the middle of the data."""
         # Create initial data
         data = np.zeros((5, 5, 5))
 
@@ -228,12 +186,81 @@ def hole() -> NDArray:
 
         return data
 
+    @staticmethod
+    def kill_you() -> NDArray:
+        """Create a complex test case combining several simpler cases."""
+        # Create initial data
+        data = np.zeros((10, 10, 10))
+        data[3:-3, 3:-3, 3:-3] = 1
+
+        # Case 1 - simple extrusion
+        data[2, 4, 4] = 1
+        data[2, 5, 5] = 1
+
+        # Case 2 - double extrusion
+        data[4, 4, 2] = 1
+        data[4, 4, 1] = 1
+        data[5, 5, 2] = 1
+        data[5, 5, 1] = 1
+
+        # Case 3 - triple extrusion
+        data[4, 7, 4] = 1
+        data[4, 8, 4] = 1
+        data[4, 9, 4] = 1
+        data[5, 7, 5] = 1
+        data[5, 8, 5] = 1
+        data[5, 9, 5] = 1
+
+        # Case 4 - stairs
+        data[4, 4, 7] = 1
+        data[4, 4, 8] = 1
+        data[4, 5, 7] = 1
+        data[5, 5, 7] = 1
+        data[5, 5, 8] = 1
+        data[5, 5, 9] = 1
+
+        # Case 5 - ceiling
+        data[4, 1, 4] = 1
+        data[4, 1, 5] = 1
+        data[4, 2, 4] = 1
+        data[5, 1, 4] = 1
+        data[5, 1, 5] = 1
+        data[5, 2, 5] = 1
+
+        # Case 6 - checkboard
+        data[7, 5, 4] = 1
+        data[7, 4, 5] = 1
+        data[8, 4, 4] = 1
+        data[8, 3, 4] = 1
+        data[8, 3, 3] = 1
+        data[7, 3, 3] = 1
+        data[8, 5, 5] = 1
+        data[8, 6, 5] = 1
+        data[8, 6, 6] = 1
+        data[7, 6, 6] = 1
+
+        return data
+
     @staticmethod
     def random(*, size: int = 16, seed: int | None = None) -> NDArray:
+        """Create a random test case.
+
+        Parameters
+        ----------
+        size : int, optional
+            The size of the data cube, by default 16
+        seed : int | None, optional
+            The random seed, by default None
+
+        Returns
+        -------
+        NDArray
+            The random test case.
+        """
         rng = np.random.default_rng(seed)
         data = rng.integers(0, 2, (size, size, size), dtype=np.uint8)
         data = cc3d.largest_k(data, k=1, connectivity=6, binary_image=True).astype(
-            np.uint8
+            np.uint8,
         )
         data = fill_voids.fill(data, in_place=True)
         data = fill_voids.fill(data, in_place=True)

edge_mender/tests/test_data_factory.py

@@ -0,0 +1,29 @@
+import numpy as np
+import pytest
+from numpy.typing import NDArray
+
+from edge_mender.data_factory import DataFactory
+
+
+@pytest.mark.parametrize(
+    "data",
+    [
+        DataFactory.simple_extrusion(),
+        DataFactory.double_extrusion(),
+        DataFactory.triple_extrusion(),
+        DataFactory.stairs(),
+        DataFactory.ceiling(),
+        DataFactory.double_tower_ceiling(),
+        DataFactory.hanging_points(),
+        DataFactory.checkerboard(),
+        DataFactory.hole(),
+        DataFactory.kill_you(),
+        DataFactory.random(seed=0),
+    ],
+)
+def test_data_factory(data: NDArray) -> None:
+    """Test that the data factory creates the expected data."""
+    assert data is not None
+    assert isinstance(data, np.ndarray)
+    assert data.ndim == 3  # noqa: PLR2004
+    assert np.any(data == 1)

edge_mender/tests/test_edge_mender.py

@@ -23,6 +23,7 @@
         DataFactory.hanging_points(),
         DataFactory.checkerboard(),
         DataFactory.hole(),
+        DataFactory.kill_you(),
     ],
 )
 def test_validate(data: NDArray, spacing: tuple[float, float, float]) -> None:
@@ -192,7 +193,8 @@ def test_find_non_manifold_edges(
         DataFactory.checkerboard(),
         # TODO: This test case fails due to a bug with SurfaceNets from VTK
         # https://gitlab.kitware.com/vtk/vtk/-/issues/19156, fixed, but not released yet
-        # DataFactory.hole(.)
+        # DataFactory.hole(),  # noqa: ERA001
+        DataFactory.kill_you(),
     ],
 )
 def test_repair(data: NDArray) -> None: