TST: Make assert_warns an optional contextmanager

1 files changed, 15 insertions, 0 deletions

diff --git a/numpy/testing/tests/test_utils.py b/numpy/testing/tests/test_utils.py
index 92a00f712..46c7fde5b 100644
--- a/numpy/testing/tests/test_utils.py
+++ b/numpy/testing/tests/test_utils.py
@@ -465,6 +465,21 @@ class TestWarns(unittest.TestCase):
         assert_equal(before_filters, after_filters,
                      "assert_warns does not preserver warnings state")
 
+    def test_context_manager(self):
+
+        before_filters = sys.modules['warnings'].filters[:]
+        with assert_warns(UserWarning):
+            warnings.warn("yo")
+        after_filters = sys.modules['warnings'].filters
+
+        def no_warnings():
+            with assert_no_warnings():
+                warnings.warn("yo")
+
+        assert_raises(AssertionError, no_warnings)
+        assert_equal(before_filters, after_filters,
+                     "assert_warns does not preserver warnings state")
+
     def test_warn_wrong_warning(self):
         def f():
             warnings.warn("yo", DeprecationWarning)