Add checks / tests for AFIDs coord array shape

afids_utils/io.py

@@ -90,6 +90,17 @@ def afids_to_fcsv(
         )
         fcsv = list(reader)
 
+    # Check to make sure shape of AFIDs array matches expected template
+    if afid_coords.shape[0] != len(fcsv):
+        raise TypeError(
+            f"Expected {len(fcsv)} AFIDs, but received {afid_coords.shape[0]}."
+        )
+    if afid_coords.shape[1] != 3:
+        raise TypeError(
+            "Expected 3 spatial dimensions (x, y, z),"
+            f"but received {afid_coords.shape[1]}."
+        )
+
     # Loop over fiducials and update with fiducial spatial coordinates
     for idx, row in enumerate(fcsv):
         row["x"] = afid_coords[idx][0]

afids_utils/tests/strategies.py

@@ -12,10 +12,23 @@ def afid_coords(
     min_value: float = -50.0,
     max_value: float = 50.0,
     width: int = 16,
+    bad_range: bool = False,
+    bad_dims: bool = False,
 ) -> NDArray[np.single]:
+    # Set (in)valid dimensions for array containing AFID coords
+    num_afids, spatial_dims = 32, 3
+    if bad_range:
+        num_afids = draw(
+            st.integers(min_value=0, max_value=100).filter(lambda x: x != 32)
+        )
+    if bad_dims:
+        spatial_dims = draw(
+            st.integers(min_value=0, max_value=10).filter(lambda x: x != 3)
+        )
+
     return draw(
         arrays(
-            shape=(32, 3),
+            shape=(num_afids, spatial_dims),
             dtype=np.single,
             elements=st.floats(
                 min_value=min_value, max_value=max_value, width=width

afids_utils/tests/test_io.py

@@ -2,7 +2,7 @@
 
 import csv
 import tempfile
-from os import PathLike, remove
+from os import PathLike
 from pathlib import Path
 
 import numpy as np
@@ -64,47 +64,62 @@ def test_invalid_template(self, afids_coords: NDArray[np.single]) -> None:
                 "/invalid/fcsv/path",
             )
 
+    @given(afids_coords=afid_coords(bad_range=True))
+    def test_invalid_num_afids(self, afids_coords: NDArray[np.single]) -> None:
+        with tempfile.NamedTemporaryFile(
+            mode="w", prefix="sub-test_desc-", suffix="_afids.fcsv"
+        ) as out_fcsv_file:
+            with pytest.raises(TypeError) as err:
+                afids_to_fcsv(afids_coords, out_fcsv_file)
+
+            assert "AFIDs, but received" in str(err.value)
+
+    @given(afids_coords=afid_coords(bad_dims=True))
+    def test_invalid_afids_dims(
+        self, afids_coords: NDArray[np.single]
+    ) -> None:
+        with tempfile.NamedTemporaryFile(
+            mode="w", prefix="sub-test_desc-", suffix="_afids.fcsv"
+        ) as out_fcsv_file:
+            with pytest.raises(TypeError) as err:
+                afids_to_fcsv(afids_coords, out_fcsv_file)
+
+            assert "Expected 3 spatial dimensions" in str(err.value)
+
     @given(afids_coords=afid_coords())
     @settings(
         suppress_health_check=[HealthCheck.function_scoped_fixture],
     )
     def test_write_fcsv(
         self, afids_coords: NDArray[np.single], valid_fcsv_file: PathLike[str]
     ) -> None:
-        out_fcsv_file = tempfile.NamedTemporaryFile(
-            mode="w", delete=False, prefix="sub-test_afids.fcsv"
-        )
-        out_fcsv_path = Path(out_fcsv_file.name)
-
-        afids_to_fcsv(afids_coords, out_fcsv_path)
-
-        # Check file was created
-        assert out_fcsv_path.exists()
-
-        # Load files
-        with open(
-            valid_fcsv_file, "r", encoding="utf-8", newline=""
-        ) as template_fcsv_file:
-            template_header = [template_fcsv_file.readline() for _ in range(3)]
-
-        with open(
-            out_fcsv_path, "r", encoding="utf-8", newline=""
-        ) as output_fcsv_file:
-            output_header = [output_fcsv_file.readline() for _ in range(3)]
-            reader = csv.DictReader(
-                output_fcsv_file, fieldnames=FCSV_FIELDNAMES
-            )
-            output_fcsv = list(reader)
-
-        # Check header
-        assert output_header == template_header
-        # Check contents
-        for idx, row in enumerate(output_fcsv):
-            assert (row["x"], row["y"], row["z"]) == (
-                str(afids_coords[idx][0]),
-                str(afids_coords[idx][1]),
-                str(afids_coords[idx][2]),
-            )
-
-        # Delete temporary file
-        remove(out_fcsv_path)
+        with tempfile.NamedTemporaryFile(
+            mode="w", prefix="sub-test_desc-", suffix="_afids.fcsv"
+        ) as out_fcsv_file:
+            # Create and check output file
+            afids_to_fcsv(afids_coords, out_fcsv_file.name)
+
+            # Load files
+            with open(
+                valid_fcsv_file, "r", encoding="utf-8", newline=""
+            ) as template_fcsv_file, open(
+                out_fcsv_file.name, "r", encoding="utf-8", newline=""
+            ) as output_fcsv_file:
+                template_header = [
+                    template_fcsv_file.readline() for _ in range(3)
+                ]
+                output_header = [output_fcsv_file.readline() for _ in range(3)]
+                reader = csv.DictReader(
+                    output_fcsv_file, fieldnames=FCSV_FIELDNAMES
+                )
+                output_fcsv = list(reader)
+
+            # Check header
+            assert output_header == template_header
+            # Check contents
+            for idx, row in enumerate(output_fcsv):
+                assert (row["x"], row["y"], row["z"]) == (
+                    str(afids_coords[idx][0]),
+                    str(afids_coords[idx][1]),
+                    str(afids_coords[idx][2]),
+                )