100% test coverage: add test for unsupported correction type

src/correctionlib_gradients/_base.py

@@ -14,21 +14,29 @@
 Value: TypeAlias = float | np.ndarray | jax.Array
 
 
-def apply_ast(ast: schema.Content, correction_name: str, _inputs: dict[str, Value]) -> Value:
+def apply_ast(ast: schema.Content, _inputs: dict[str, Value]) -> Value:
     match ast:
         case float(x):
             return x
+        case _:  # pragma: no cover
+            msg = "Unsupported type of node in the computation graph. This should never happen."
+            raise RuntimeError(msg)
+
+
+def assert_supported(c: schema.Correction, name: str) -> None:
+    match c.data:
+        case float():
+            return
         case _:
-            msg = (
-                f"Cannot compute gradients of correction '{correction_name}': "
-                f"it contains the unsupported operation type '{type(ast).__name__}'"
-            )
-            raise NotImplementedError(msg)
+            msg = f"Correction '{name}' contains the unsupported operation type '{type(c.data).__name__}'"
+            raise ValueError(msg)
 
 
 class CorrectionWithGradient:
     def __init__(self, c: schema.Correction):
-        self._evaluator = partial(apply_ast, c.data, c.name)
+        assert_supported(c, c.name)
+
+        self._evaluator = partial(apply_ast, c.data)
         self._input_names = [v.name for v in c.inputs]
         self._name = c.name
 

tests/test_base.py

@@ -21,7 +21,19 @@
         inputs=[schemav2.Variable(name="x", type="real")],
         output=schemav2.Variable(name="a scale", type="real"),
         data=1.234,
-    )
+    ),
+    # this type of correction is unsupported
+    "categorical": schemav2.Correction(
+        name="categorical",
+        version=2,
+        inputs=[schemav2.Variable(name="c", type="int")],
+        output=schemav2.Variable(name="a scale", type="real"),
+        data=schemav2.Category(
+            nodetype="category",
+            input="x",
+            content=[schemav2.CategoryItem(key=0, value=1.234)],
+        ),
+    ),
 }
 
 
@@ -41,6 +53,11 @@ def test_missing_input():
         cg.eval_dict({})
 
 
+def test_unsupported_correction():
+    with pytest.raises(ValueError, match="Correction 'categorical' contains the unsupported operation type 'Category'"):
+        CorrectionWithGradient(schemas["categorical"])
+
+
 @pytest.mark.parametrize("jit", [False, True])
 def test_evaluate_scale(jit):
     cg = CorrectionWithGradient(schemas["scale"])