diff --git a/pandas/tests/util/test_assert_produces_warning.py b/pandas/tests/util/test_assert_produces_warning.py new file mode 100644 index 0000000000000..e125f01a494e7 --- /dev/null +++ b/pandas/tests/util/test_assert_produces_warning.py @@ -0,0 +1,22 @@ +import warnings + +import pytest + +import pandas.util.testing as tm + + +def f(): + warnings.warn('f1', FutureWarning) + warnings.warn('f2', RuntimeWarning) + + +@pytest.mark.filterwarnings('ignore:f1:FutureWarning') +def test_assert_produces_warning_honors_filter(): + # raise by default + with pytest.raises(AssertionError): + with tm.assert_produces_warning(RuntimeWarning): + f() + + with tm.assert_produces_warning(RuntimeWarning, + raise_on_extra_warnings=False): + f() diff --git a/pandas/util/testing.py b/pandas/util/testing.py index 6e88cd7f72dcd..0481e88bb9a05 100644 --- a/pandas/util/testing.py +++ b/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, + raise_on_extra_warnings=True): """ Context manager for running code expected to either raise a specific warning, or not raise any warnings. Verifies that the code raises the @@ -2584,7 +2585,7 @@ def assert_produces_warning(expected_warning=Warning, filter_level="always", The type of Exception raised. ``exception.Warning`` is the base class for all warnings. To check that no warning is returned, specify ``False`` or ``None``. - filter_level : str, default "always" + filter_level : str or None, default "always" Specifies whether warnings are ignored, displayed, or turned into errors. Valid values are: @@ -2608,6 +2609,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. + raise_on_extra_warnings : bool, default True + Whether extra warnings not of the type `expected_warning` should + cause the test to fail. Examples -------- @@ -2676,8 +2680,10 @@ class for all warnings. To check that no warning is returned, msg = "Did not see expected warning of class {name!r}.".format( name=expected_warning.__name__) assert saw_warning, msg - assert not extra_warnings, ("Caused unexpected warning(s): {extra!r}." - ).format(extra=extra_warnings) + if raise_on_extra_warnings and extra_warnings: + raise AssertionError( + "Caused unexpected warning(s): {!r}.".format(extra_warnings) + ) class RNGContext(object):