migrate assert_warns_message to imblearn.utils.testing.warns

Author: Joan Massich

imblearn/datasets/tests/test_imbalance.py

@@ -11,7 +11,8 @@
 import numpy as np
 
 from sklearn.datasets import load_iris
-from sklearn.utils.testing import assert_warns_message
+from imblearn.utils.testing import warns
+
 from pytest import raises
 
 from imblearn.datasets import make_imbalance
@@ -40,12 +41,12 @@ def test_make_imbalance_error():
 
 # FIXME: to be removed in 0.4 due to deprecation
 def test_make_imbalance_float():
-    X_, y_ = assert_warns_message(DeprecationWarning,
-                                  "'min_c_' is deprecated in 0.2",
-                                  make_imbalance, X, Y, ratio=0.5, min_c_=1)
-    X_, y_ = assert_warns_message(DeprecationWarning,
-                                  "'ratio' being a float is deprecated",
-                                  make_imbalance, X, Y, ratio=0.5, min_c_=1)
+    with warns(DeprecationWarning, match="deprecated in 0.2"):
+        X_, y_ = make_imbalance(X, Y, ratio=0.5, min_c_=1)
+
+    with warns(DeprecationWarning, match="'ratio' being a float"):
+        X_, y_ = make_imbalance(X, Y, ratio=0.5, min_c_=1)
+
     assert Counter(y_) == {0: 50, 1: 25, 2: 50}
     # resample without using min_c_
     X_, y_ = make_imbalance(X_, y_, ratio=0.25, min_c_=None)

imblearn/metrics/tests/test_classification.py

@@ -18,7 +18,7 @@
 from sklearn.utils.validation import check_random_state
 from sklearn.utils.testing import assert_allclose, assert_array_equal
 from sklearn.utils.testing import assert_no_warnings
-from sklearn.utils.testing import assert_warns_message, ignore_warnings
+from sklearn.utils.testing import ignore_warnings
 from sklearn.metrics import accuracy_score, average_precision_score
 from sklearn.metrics import brier_score_loss, cohen_kappa_score
 from sklearn.metrics import jaccard_similarity_score, precision_score
@@ -32,6 +32,8 @@
 from imblearn.metrics import classification_report_imbalanced
 
 from pytest import approx, raises
+from imblearn.utils.testing import warns
+
 
 RND_SEED = 42
 R_TOL = 1e-2
@@ -196,15 +198,10 @@ def test_sensitivity_specificity_support_errors():
 
 def test_sensitivity_specificity_unused_pos_label():
     # but average != 'binary'; even if data is binary
-    assert_warns_message(
-        UserWarning,
-        "Note that pos_label (set to 2) is "
-        "ignored when average != 'binary' (got 'macro'). You "
-        "may use labels=[pos_label] to specify a single "
-        "positive class.",
-        sensitivity_specificity_support, [1, 2, 1], [1, 2, 2],
-        pos_label=2,
-        average='macro')
+    with warns(UserWarning, "use labels=\[pos_label\] to specify a single"):
+        sensitivity_specificity_support([1, 2, 1], [1, 2, 2],
+                                        pos_label=2,
+                                        average='macro')
 
 
 def test_geometric_mean_support_binary():

imblearn/utils/tests/test_deprecation.py

@@ -3,9 +3,8 @@
 # Authors: Guillaume Lemaitre <g.lemaitre58@gmail.com>
 # License: MIT
 
-from sklearn.utils.testing import assert_warns_message
-
 from imblearn.utils.deprecation import deprecate_parameter
+from imblearn.utils.testing import warns
 
 
 class Sampler(object):
@@ -15,7 +14,7 @@ def __init__(self):
 
 
 def test_deprecate_parameter():
-    assert_warns_message(DeprecationWarning, "is deprecated from",
-                         deprecate_parameter, Sampler(), '0.2', 'a')
-    assert_warns_message(DeprecationWarning, "Use 'b' instead.",
-                         deprecate_parameter, Sampler(), '0.2', 'a', 'b')
+    with warns(DeprecationWarning, match="is deprecated from"):
+        deprecate_parameter(Sampler(), '0.2', 'a')
+    with warns(DeprecationWarning, match="Use 'b' instead."):
+        deprecate_parameter(Sampler(), '0.2', 'a', 'b')

imblearn/utils/tests/test_validation.py

@@ -10,7 +10,7 @@
 from sklearn.neighbors.base import KNeighborsMixin
 from sklearn.neighbors import NearestNeighbors
 
-from sklearn.utils.testing import assert_warns_message
+from imblearn.utils.testing import warns
 from pytest import raises
 
 from imblearn.utils import check_neighbors_object
@@ -126,10 +126,11 @@ def test_ratio_dict_over_sampling():
     ratio_ = check_ratio(ratio, y, 'over-sampling')
     assert ratio_ == {1: 20, 2: 0, 3: 45}
     ratio = {1: 70, 2: 140, 3: 70}
-    assert_warns_message(UserWarning, "After over-sampling, the number of"
-                         " samples (140) in class 2 will be larger than the"
-                         " number of samples in the majority class (class #2"
-                         " -> 100)", check_ratio, ratio, y, 'over-sampling')
+    expected_msg = ("After over-sampling, the number of samples \(140\) in"
+                    " class 2 will be larger than the number of samples in the"
+                    " majority class \(class #2 -> 100\)")
+    with warns(UserWarning, expected_msg):
+        check_ratio(ratio, y, 'over-sampling')
 
 
 def test_ratio_dict_under_sampling():