add warns context manger that mimics raises

Author: Joan Massich

imblearn/utils/testing.py

@@ -14,6 +14,11 @@
 
 from sklearn.base import BaseEstimator
 
+from pytest import warns as _warns
+from contextlib import contextmanager
+from re import compile
+
+
 # meta-estimators need another estimator to be instantiated.
 META_ESTIMATORS = []
 # estimators that there is no way to default-construct sensibly
@@ -120,3 +125,72 @@ def is_abstract(c):
     # itemgetter is used to ensure the sort does not extend to the 2nd item of
     # the tuple
     return sorted(set(estimators), key=itemgetter(0))
+
+
+@contextmanager
+def warns(expected_warning, message=None, match=None):
+    """
+    Assert that a code block/function call warns ``expected_warning``
+    and raise a failure exception otherwise.
+
+    If using Python 2.7 or above, you may use this function as a
+    context manager::
+
+        >>> import warnings
+        >>> with warns(RuntimeWarning):
+        ...    warnings.warn("my runtime warning", RuntimeWarning)
+
+        >>> with warns(RuntimeWarning):
+        ...    pass
+        Traceback (most recent call last):
+          ...
+        Failed: DID NOT WARN. No warnings of type (<class 'RuntimeError'>,) was emitted. The list of emitted warnings is: [].
+
+        >>> with warns(RuntimeWarning):
+        ...    warnings.warn(UserWarning)
+        Traceback (most recent call last):
+          ...
+        Failed: DID NOT WARN. No warnings of type (<class 'RuntimeError'>,) was emitted. The list of emitted warnings is: [UserWarning(<class 'UserWarning'>,)].
+
+
+    In the context manager form you may use the keyword argument
+    ``message`` to specify a custom failure message::
+
+        >>> import warnings
+        >>> with warns(RuntimeWarning, message="my runtime warning"):
+        ...    warnings.warn("different message", RuntimeWarning)
+        Traceback (most recent call last):
+          ...
+        AssertionError: "different message" is different from "my runtime warning"
+
+    Or you can use the keyword argument ``match`` to assert that the
+    exception matches a text or regex::
+
+        >>> import warnings
+        >>> with warns(UserWarning, match='must be 0 or None'):
+        ...     warnings.warn("value must be 0 or None", UserWarning)
+
+        >>> with warns(UserWarning, match=r'must be \d+$'):
+        ...     warnings.warn("value must be 42", UserWarning)
+
+        >>> with warns(UserWarning, match=r'must be \d+$'):
+        ...     warnings.warn("this is not here", UserWarning)
+        Traceback (most recent call last):
+          ...
+        AssertionError: "must be \d+$" pattern not found in "this is not here"
+
+    """
+    with _warns(expected_warning) as record:
+        yield
+
+    wrn_msg = str(record[0].message)
+    if message is not None:
+        # message matching has priority over regex match
+        assert_msg = '"{}" is different from "{}"'.format(wrn_msg, message)
+        assert wrn_msg == message, assert_msg
+    elif match is not None:
+        regex = compile(match)
+        assert_msg = '"{}" pattern not found in "{}"'.format(match, wrn_msg)
+        assert regex.search(wrn_msg) is not None, assert_msg
+    else:
+        pass

imblearn/utils/tests/test_testing.py

@@ -8,6 +8,9 @@
 from imblearn.base import SamplerMixin
 from imblearn.utils.testing import all_estimators
 
+from imblearn.utils.testing import warns
+from pytest import raises
+
 
 def test_all_estimators():
     # check if the filtering is working with a list or a single string
@@ -23,3 +26,21 @@ def test_all_estimators():
     type_filter = 'rnd'
     with raises(ValueError, match="Parameter type_filter must be 'sampler'"):
         all_estimators(type_filter=type_filter)
+
+
+def test_warns():
+    import warnings
+
+    with warns(RuntimeWarning, message='my runtime warning'):
+        warnings.warn("my runtime warning", RuntimeWarning)
+
+    with warns(UserWarning, match=r'must be \d+$'):
+        warnings.warn("value must be 42", UserWarning)
+
+    with raises(AssertionError, match='is different from'):
+        with warns(RuntimeWarning, message="my runtime warning"):
+            warnings.warn("different message", RuntimeWarning)
+
+    with raises(AssertionError, match='pattern not found'):
+        with warns(UserWarning, match=r'must be \d+$'):
+            warnings.warn("this is not here", UserWarning)