Issue #4444: Allow assertRaises() to be used as a context handler.

1 files changed, 30 insertions, 8 deletions

diff --git a/Lib/unittest.py b/Lib/unittest.py
index 09c6ca9..4538e30 100644
--- a/Lib/unittest.py
+++ b/Lib/unittest.py
@@ -174,6 +174,25 @@ class TestResult:
                (_strclass(self.__class__), self.testsRun, len(self.errors),
                 len(self.failures))
 
+class AssertRaisesContext:
+    def __init__(self, expected, test_case):
+        self.expected = expected
+        self.failureException = test_case.failureException
+    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)
+            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.
 
@@ -324,22 +343,25 @@ class TestCase:
         """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)
+        if callableObj is None:
+            return context
+        with context:
             callableObj(*args, **kwargs)
-        except excClass:
-            return
-        else:
-            if hasattr(excClass,'__name__'): excName = excClass.__name__
-            else: excName = str(excClass)
-            raise self.failureException, "%s not raised" % excName
 
     def failUnlessEqual(self, first, second, msg=None):
         """Fail if the two objects are unequal as determined by the '=='