TST Replace pytest.warns(None) in test_omp.py (#22886)

richardt94 · web-flow · commit aee564c544e3 · 2022-03-17T20:13:51.000-04:00
diff --git a/sklearn/linear_model/tests/test_omp.py b/sklearn/linear_model/tests/test_omp.py
@@ -3,6 +3,7 @@
 
 import numpy as np
 import pytest
+import warnings
 
 from sklearn.utils._testing import assert_allclose
 from sklearn.utils._testing import assert_array_equal
@@ -53,11 +54,11 @@ def test_assure_warning_when_normalize(OmpModel, normalize, n_warnings):
     y = rng.rand(n_samples)
 
     model = OmpModel(normalize=normalize)
-    with pytest.warns(None) as record:
+    with warnings.catch_warnings(record=True) as rec:
+        warnings.simplefilter("always", FutureWarning)
         model.fit(X, y)
 
-    record = [r for r in record if r.category == FutureWarning]
-    assert len(record) == n_warnings
+    assert len([w.message for w in rec]) == n_warnings
 
 
 def test_correct_shapes():
