try to improve code coverage

ornlneutronimaging · Feb 16, 2024 · 750af05 · 750af05
1 parent 4d232f2
commit 750af05
Showing 1 changed file with 16 additions and 0 deletions.
diff --git a/tests/unit/backend/corrections/test_beam_hardening.py b/tests/unit/backend/corrections/test_beam_hardening.py
@@ -39,6 +39,22 @@ def fake_beam_hardening_image() -> np.ndarray:
 
 def test_beam_hardening_correction(fake_beam_hardening_image):
     """Test beam hardening correction"""
+    # test with incorrect array dim
+    with pytest.raises(ValueError):
+        beam_hardening_correction(
+            arrays=np.array([0, 1, 2, 3]),
+            q=0.005,
+            n=20.0,
+            opt=True,
+        )
+    # test with single image
+    corrected_image = beam_hardening_correction(
+        arrays=fake_beam_hardening_image[0],
+        q=0.005,
+        n=20.0,
+        opt=True,
+    )
+    assert corrected_image.shape == (255, 255)
     # test the correction
     corrected_image = beam_hardening_correction(
         arrays=fake_beam_hardening_image,