MNT Replace pytest.warns(None) in linear_models tests (#23138)

thomasjpfan · web-flow · commit 3d096378ff8c · 2022-04-19T15:53:55.000+02:00
diff --git a/sklearn/linear_model/tests/test_base.py b/sklearn/linear_model/tests/test_base.py
@@ -184,14 +184,15 @@ def test_deprecate_normalize(normalize, default):
             expected = None
             warning_msg = []
 
-    with pytest.warns(expected) as record:
-        _normalize = _deprecate_normalize(normalize, default, "estimator")
-    assert _normalize == output
-
-    n_warnings = 0 if expected is None else 1
-    assert len(record) == n_warnings
-    if n_warnings:
+    if expected is None:
+        with warnings.catch_warnings():
+            warnings.simplefilter("error", FutureWarning)
+            _normalize = _deprecate_normalize(normalize, default, "estimator")
+    else:
+        with pytest.warns(expected) as record:
+            _normalize = _deprecate_normalize(normalize, default, "estimator")
         assert all([warning in str(record[0].message) for warning in warning_msg])
+    assert _normalize == output
 
 
 def test_linear_regression_sparse(random_state=0):
diff --git a/sklearn/linear_model/tests/test_common.py b/sklearn/linear_model/tests/test_common.py
@@ -5,6 +5,7 @@
 import pytest
 
 import sys
+import warnings
 import numpy as np
 
 from sklearn.base import is_classifier
@@ -56,6 +57,12 @@ def test_linear_model_normalize_deprecation_message(
         y = np.sign(y)
 
     model = estimator(normalize=normalize)
+    if warning_category is None:
+        with warnings.catch_warnings():
+            warnings.simplefilter("error", FutureWarning)
+            model.fit(X, y)
+        return
+
     with pytest.warns(warning_category) as record:
         model.fit(X, y)
     # Filter record in case other unrelated warnings are raised
@@ -67,6 +74,5 @@ def test_linear_model_normalize_deprecation_message(
             msg += "\n"
         raise AssertionError(msg)
     wanted = [r for r in record if r.category == warning_category]
-    if warning_category is not None:
-        assert "'normalize' was deprecated" in str(wanted[0].message)
+    assert "'normalize' was deprecated" in str(wanted[0].message)
     assert len(wanted) == n_warnings
