Merged revisions 67979 via svnmerge from

pitrou · pitrou · commit 5acd41e0a166 · 2008-12-28T14:29:00.000Z
svn+ssh://pythondev@svn.python.org/python/trunk ........ r67979 | antoine.pitrou | 2008-12-28 15:09:36 +0100 (dim., 28 déc. 2008) | 3 lines Issue #4444: Allow assertRaises() to be used as a context handler. ........
diff --git a/Doc/library/unittest.rst b/Doc/library/unittest.rst
@@ -609,15 +609,20 @@ failures.
    equal, the test will fail with the explanation given by *msg*, or :const:`None`.
 
 
-.. method:: TestCase.assertRaises(exception, callable, ...)
-            TestCase.failUnlessRaises(exception, callable, ...)
+.. method:: TestCase.assertRaises(exception[, callable, ...])
+            TestCase.failUnlessRaises(exception[, callable, ...])
 
    Test that an exception is raised when *callable* is called with any positional
    or keyword arguments that are also passed to :meth:`assertRaises`.  The test
    passes if *exception* is raised, is an error if another exception is raised, or
    fails if no exception is raised.  To catch any of a group of exceptions, a tuple
    containing the exception classes may be passed as *exception*.
 
+   If *callable* is omitted or None, returns a context manager so that the code
+   under test can be written inline rather than as a function::
+
+        with self.failUnlessRaises(some_error_class):
+            do_something()
 
 .. method:: TestCase.failIf(expr[, msg])
             TestCase.assertFalse(expr[, msg])
diff --git a/Lib/test/test_unittest.py b/Lib/test/test_unittest.py
@@ -2284,6 +2284,43 @@ def test_AlmostEqual(self):
         self.assertRaises(AssertionError,
                           self.failIfAlmostEqual, 0, .1+.1j, places=0)
 
+    def test_assertRaises(self):
+        def _raise(e):
+            raise e
+        self.assertRaises(KeyError, _raise, KeyError)
+        self.assertRaises(KeyError, _raise, KeyError("key"))
+        try:
+            self.assertRaises(KeyError, lambda: None)
+        except AssertionError as e:
+            self.assert_("KeyError not raised" in str(e), str(e))
+        else:
+            self.fail("assertRaises() didn't fail")
+        try:
+            self.assertRaises(KeyError, _raise, ValueError)
+        except ValueError:
+            pass
+        else:
+            self.fail("assertRaises() didn't let exception pass through")
+        with self.assertRaises(KeyError):
+            raise KeyError
+        with self.assertRaises(KeyError):
+            raise KeyError("key")
+        try:
+            with self.assertRaises(KeyError):
+                pass
+        except AssertionError as e:
+            self.assert_("KeyError not raised" in str(e), str(e))
+        else:
+            self.fail("assertRaises() didn't fail")
+        try:
+            with self.assertRaises(KeyError):
+                raise ValueError
+        except ValueError:
+            pass
+        else:
+            self.fail("assertRaises() didn't let exception pass through")
+
+
 ######################################################################
 ## Main
 ######################################################################
diff --git a/Lib/unittest.py b/Lib/unittest.py
@@ -149,6 +149,36 @@ def __repr__(self):
                (_strclass(self.__class__), self.testsRun, len(self.errors),
                 len(self.failures))
 
+class AssertRaisesContext:
+    def __init__(self, expected, test_case, callable_obj=None):
+        self.expected = expected
+        self.failureException = test_case.failureException
+        if callable_obj is not None:
+            try:
+                self.obj_name = callable_obj.__name__
+            except AttributeError:
+                self.obj_name = str(callable_obj)
+        else:
+            self.obj_name = None
+    def __enter__(self):
+        pass
+    def __exit__(self, exc_type, exc_value, traceback):
+        if exc_type is None:
+            try:
+                exc_name = self.expected.__name__
+            except AttributeError:
+                exc_name = str(self.expected)
+            if self.obj_name:
+                raise self.failureException("{0} not raised by {1}"
+                    .format(exc_name, self.obj_name))
+            else:
+                raise self.failureException("{0} not raised"
+                    .format(exc_name))
+        if issubclass(exc_type, self.expected):
+            return True
+        # Let unexpected exceptions skip through
+        return False
+
 class TestCase:
     """A class whose instances are single test cases.
 
@@ -299,23 +329,25 @@ def failUnless(self, expr, msg=None):
         """Fail the test unless the expression is true."""
         if not expr: raise self.failureException(msg)
 
-    def failUnlessRaises(self, excClass, callableObj, *args, **kwargs):
+    def failUnlessRaises(self, excClass, callableObj=None, *args, **kwargs):
         """Fail unless an exception of class excClass is thrown
            by callableObj when invoked with arguments args and keyword
            arguments kwargs. If a different type of exception is
            thrown, it will not be caught, and the test case will be
            deemed to have suffered an error, exactly as for an
            unexpected exception.
+
+           If called with callableObj omitted or None, will return a
+           context object used like this::
+
+                with self.failUnlessRaises(some_error_class):
+                    do_something()
         """
-        try:
+        context = AssertRaisesContext(excClass, self, callableObj)
+        if callableObj is None:
+            return context
+        with context:
             callableObj(*args, **kwargs)
-        except excClass:
-            return
-        else:
-            excName = str(getattr(excClass, '__name__', excClass))
-            objName = str(getattr(callableObj, '__name__', callableObj))
-            raise self.failureException("%s not raised by %s" % (excName,
-                                                                  objName))
 
     def failUnlessEqual(self, first, second, msg=None):
         """Fail if the two objects are unequal as determined by the '=='
diff --git a/Misc/NEWS b/Misc/NEWS
@@ -64,6 +64,9 @@ Core and Builtins
 Library
 -------
 
+- Issue #4444: Allow assertRaises() to be used as a context handler, so that
+  the code under test can be written inline if more practical.
+
 - Issue #4756: zipfile.is_zipfile() now supports file-like objects. Patch by
   Gabriel Genellina.
 
