ENH: add an assert_warns testing utility.

1 files changed, 34 insertions, 0 deletions

diff --git a/numpy/testing/tests/test_utils.py b/numpy/testing/tests/test_utils.py
index ab314a703..2d22789ff 100644
--- a/numpy/testing/tests/test_utils.py
+++ b/numpy/testing/tests/test_utils.py
@@ -1,3 +1,6 @@
+import warnings
+import sys
+
 import numpy as np
 from numpy.testing import *
 import unittest
@@ -301,6 +304,37 @@ class TestRaises(unittest.TestCase):
         else:
             raise AssertionError("should have raised an AssertionError")
 
+class TestWarns(unittest.TestCase):
+    def test_warn(self):
+        def f():
+            warnings.warn("yo")
+
+        before_filters = sys.modules['warnings'].filters[:]
+        assert_warns(UserWarning, f)
+        after_filters = sys.modules['warnings'].filters
+
+        # Check that the warnings state is unchanged
+        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)
+
+        failed = False
+        filters = sys.modules['warnings'].filters[:]
+        try:
+            # Should raise an AssertionError
+            assert_warns(UserWarning, f)
+            failed = True
+        except AssertionError:
+            pass
+        finally:
+            sys.modules['warnings'].filters = filters
+
+        if failed:
+            raise AssertionError("wrong warning caught by assert_warn")
+
 class TestArrayAlmostEqualNulp(unittest.TestCase):
     def test_simple(self):
         dev = np.random.randn(10)