From 6cbc7e796b6a02475698dddb7272b1743e9703d6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jon=20Haitz=20Legarreta=20Gorro=C3=B1o?=
 <jon.haitz.legarreta@gmail.com>
Date: Sat, 18 Jan 2025 11:07:00 -0500
Subject: [PATCH] ENH: Test absent mask warning

Test absent mask warning: parameterize the trivial model test so that
the warning raised for the case where the model is instatiated without a
mask can be tested.

Remove the corresponding message from the `filterwarnings` list.
---
 pyproject.toml             |  2 --
 src/nifreeze/model/base.py |  6 +++++-
 test/test_model.py         | 19 ++++++++++++++++---
 3 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index aff2930..f924608 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -182,8 +182,6 @@ filterwarnings = [
   "ignore:Updating b0_threshold to.*:UserWarning",
   # scikit-learn
   "ignore:The optimal value found for dimension.*:sklearn.exceptions.ConvergenceWarning",
-  # masks
-  "ignore:No mask provided;.*:UserWarning",
 ]
 
 
diff --git a/src/nifreeze/model/base.py b/src/nifreeze/model/base.py
index bcdca3e..3d84151 100644
--- a/src/nifreeze/model/base.py
+++ b/src/nifreeze/model/base.py
@@ -28,6 +28,10 @@
 
 from nifreeze.exceptions import ModelNotFittedError
 
+mask_absence_warn_msg = (
+    "No mask provided; consider using a mask to avoid issues in model optimization."
+)
+
 
 class ModelFactory:
     """A factory for instantiating diffusion models."""
@@ -96,7 +100,7 @@ def __init__(self, mask=None, **kwargs):
 
         # Setup brain mask
         if mask is None:
-            warn("No mask provided; consider using a mask to avoid issues in model optimization.")
+            warn(mask_absence_warn_msg)
 
         self._mask = mask
 
diff --git a/test/test_model.py b/test/test_model.py
index dd901bb..b8bb155 100644
--- a/test/test_model.py
+++ b/test/test_model.py
@@ -22,6 +22,8 @@
 #
 """Unit tests exercising models."""
 
+import contextlib
+
 import numpy as np
 import pytest
 from dipy.sims.voxel import single_tensor
@@ -31,11 +33,13 @@
 from nifreeze.data.splitting import lovo_split
 from nifreeze.exceptions import ModelNotFittedError
 from nifreeze.model._dipy import GaussianProcessModel
+from nifreeze.model.base import mask_absence_warn_msg
 from nifreeze.model.dmri import DEFAULT_MAX_S0, DEFAULT_MIN_S0
 from nifreeze.testing import simulations as _sim
 
 
-def test_trivial_model():
+@pytest.mark.parametrize("use_mask", (False, True))
+def test_trivial_model(use_mask):
     """Check the implementation of the trivial B0 model."""
 
     rng = np.random.default_rng(1234)
@@ -44,7 +48,15 @@ def test_trivial_model():
     with pytest.raises(TypeError):
         model.TrivialModel()
 
-    _S0 = rng.normal(size=(2, 2, 2))
+    size = (2, 2, 2)
+    mask = None
+    if use_mask:
+        mask = np.ones(size, dtype=bool)
+        context = contextlib.nullcontext()
+    else:
+        context = pytest.warns(UserWarning, match=mask_absence_warn_msg)
+
+    _S0 = rng.normal(size=size)
 
     _clipped_S0 = np.clip(
         _S0.astype("float32") / _S0.max(),
@@ -52,7 +64,8 @@ def test_trivial_model():
         a_max=DEFAULT_MAX_S0,
     )
 
-    tmodel = model.TrivialModel(predicted=_clipped_S0)
+    with context:
+        tmodel = model.TrivialModel(mask=mask, predicted=_clipped_S0)
 
     data = None
     assert tmodel.fit(data) is None