take 2

Author: TomAugspurger

pandas/tests/util/test_assert_produces_warning.py

@@ -5,12 +5,19 @@
 import pandas.util.testing as tm
 
 
-def f():
-    warnings.warn('f1', FutureWarning)
-    warnings.warn('f2', RuntimeWarning)
+def f(a=FutureWarning, b=RuntimeWarning):
+    warnings.warn('f1', a)
+    warnings.warn('f2', b)
 
 
 @pytest.mark.filterwarnings('ignore:f1:FutureWarning')
+@pytest.mark.filterwarnings('ignore:f2:RuntimeWarning')
 def test_assert_produces_warning_honors_filter():
-    with tm.assert_produces_warning(RuntimeWarning, filter_level=None):
+    with tm.assert_produces_warning(RuntimeWarning):
         f()
+
+
+@pytest.mark.filterwarnings('ignore:f1:FutureWarning')
+def test_assert_produces_warning_message():
+    with tm.assert_produces_warning(FutureWarning, message='f2'):
+        f(FutureWarning, FutureWarning)

pandas/util/testing.py

@@ -2571,7 +2571,8 @@ def exception_matches(self, exc_type, exc_value, trace_back):
 
 @contextmanager
 def assert_produces_warning(expected_warning=Warning, filter_level="always",
-                            clear=None, check_stacklevel=True):
+                            clear=None, check_stacklevel=True,
+                            message=''):
     """
     Context manager for running code expected to either raise a specific
     warning, or not raise any warnings. Verifies that the code raises the
@@ -2609,6 +2610,9 @@ class for all warnings. To check that no warning is returned,
         If True, displays the line that called the function containing
         the warning to show were the function is called. Otherwise, the
         line that implements the function is displayed.
+    message : str, default ''
+        Use in the filter with `filter_level` and `expected_warning`
+        the control which warnings the filter applies to.
 
     Examples
     --------
@@ -2648,7 +2652,7 @@ class for all warnings. To check that no warning is returned,
 
         saw_warning = False
         if filter_level:
-            warnings.simplefilter(filter_level)
+            warnings.filterwarnings(filter_level, message, expected_warning)
         yield w
         extra_warnings = []