Merge pull request #300 from jhlegarreta/tst/fix-trivial-model-test-typing

STY: Fix `Trivial` model test typing

Author: jhlegarreta

src/nifreeze/model/base.py

@@ -27,9 +27,12 @@
 
 import numpy as np
 
-mask_absence_warn_msg = (
+MASK_ABSENCE_WARN_MSG = (
     "No mask provided; consider using a mask to avoid issues in model optimization."
 )
+"""Mask warning message."""
+PREDICTED_MAP_ERROR_MSG = "This model requires the predicted map at initialization"
+"""Oracle requirement error message."""
 UNSUPPORTED_MODEL_ERROR_MSG = "Unsupported model <{model}>."
 """Unsupported model error message"""
 
@@ -102,7 +105,7 @@ def __init__(self, dataset, **kwargs):
         self._dataset = dataset
         # Warn if mask not present
         if dataset.brainmask is None:
-            warn(mask_absence_warn_msg, stacklevel=2)
+            warn(MASK_ABSENCE_WARN_MSG, stacklevel=2)
 
     @abstractmethod
     def fit_predict(self, index: int | None = None, **kwargs) -> np.ndarray | None:
@@ -139,7 +142,7 @@ def __init__(self, dataset, predicted=None, **kwargs):
         )
 
         if self._locked_fit is None:
-            raise TypeError("This model requires the predicted map at initialization")
+            raise TypeError(PREDICTED_MAP_ERROR_MSG)
 
     def fit_predict(self, *_, **kwargs) -> np.ndarray | None:
         """Return the reference map."""

test/test_model.py

@@ -39,7 +39,11 @@
     DWI,
 )
 from nifreeze.model._dipy import GaussianProcessModel
-from nifreeze.model.base import UNSUPPORTED_MODEL_ERROR_MSG, mask_absence_warn_msg
+from nifreeze.model.base import (
+    MASK_ABSENCE_WARN_MSG,
+    PREDICTED_MAP_ERROR_MSG,
+    UNSUPPORTED_MODEL_ERROR_MSG,
+)
 from nifreeze.testing import simulations as _sim
 
 
@@ -60,6 +64,12 @@ class DummyDataset:
     pass
 
 
+class DummyDatasetNoRef:
+    def __init__(self):
+        # No reference or bzero here to trigger TrivialModel error
+        self.brainmask = np.ones((1, 1, 1, 1)).astype(bool)
+
+
 def test_base_model():
     from nifreeze.model.base import BaseModel
 
@@ -78,8 +88,8 @@ def test_trivial_model(request, use_mask):
     rng = request.node.rng
 
     # Should not allow initialization without an oracle
-    with pytest.raises(TypeError):
-        model.TrivialModel()  # type: ignore[call-arg]
+    with pytest.raises(TypeError, match=PREDICTED_MAP_ERROR_MSG):
+        model.TrivialModel(DummyDatasetNoRef())
 
     size = (2, 2, 2)
     mask = None
@@ -88,7 +98,7 @@ def test_trivial_model(request, use_mask):
         mask = np.ones(size, dtype=bool)
         context = contextlib.nullcontext()
     else:
-        context = pytest.warns(UserWarning, match=mask_absence_warn_msg)
+        context = pytest.warns(UserWarning, match=MASK_ABSENCE_WARN_MSG)
 
     _S0 = rng.normal(size=size)