Support TestCase.assertRaises being used as a context manager

LICENSE

@@ -21,6 +21,7 @@ The testtools authors are:
  * Tristan Seligmann
  * Julian Edwards
  * Jonathan Jacobs
+ * Jelmer Vernooij
 
 and are collectively referred to as "testtools developers".
 
@@ -42,7 +43,7 @@ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 SOFTWARE.
 
-Some code in testtools/run.py taken from Python's unittest module:
+Some code in testtools/run.py and testtools/testcase.py taken from Python's unittest module:
 Copyright (c) 1999-2003 Steve Purcell
 Copyright (c) 2003-2010 Python Software Foundation
 

testtools/testcase.py

@@ -436,14 +436,32 @@ def assertIsInstance(self, obj, klass, msg=None):
             matcher = IsInstance(klass)
         self.assertThat(obj, matcher, msg)
 
-    def assertRaises(self, excClass, callableObj, *args, **kwargs):
+    def assertRaises(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 the callable omitted, will return a
+        context object used like this::
+
+            with self.assertRaises(SomeException):
+                do_something()
+
+        The context manager keeps a reference to the exception as
+        the 'exception' attribute. This allows you to inspect the
+        exception after the assertion::
+
+            with self.assertRaises(SomeException) as cm:
+                do_something()
+            the_exception = cm.exception
+            self.assertEqual(the_exception.error_code, 3)
         """
+        # If callableObj is None, we're being used as a context manager
+        if callableObj is None:
+            return _AssertRaisesContext(excClass, self, msg=kwargs.get("msg"))
 
         class ReRaiseOtherTypes:
             def match(self, matchee):
@@ -1011,6 +1029,51 @@ def _id(obj):
     return _id
 
 
+class _AssertRaisesContext:
+    """A context manager to handle expected exceptions for assertRaises.
+
+    This provides compatibility with unittest's assertRaises context manager.
+    """
+
+    def __init__(self, expected, test_case, msg=None):
+        """Construct an `_AssertRaisesContext`.
+
+        :param expected: The type of exception to expect.
+        :param test_case: The TestCase instance using this context.
+        :param msg: An optional message explaining the failure.
+        """
+        self.expected = expected
+        self.test_case = test_case
+        self.msg = msg
+        self.exception = None
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if exc_type is None:
+            try:
+                if isinstance(self.expected, tuple):
+                    exc_name = "({})".format(
+                        ", ".join(e.__name__ for e in self.expected)
+                    )
+                else:
+                    exc_name = self.expected.__name__
+            except AttributeError:
+                exc_name = str(self.expected)
+            if self.msg:
+                error_msg = "{} not raised : {}".format(exc_name, self.msg)
+            else:
+                error_msg = "{} not raised".format(exc_name)
+            raise self.test_case.failureException(error_msg)
+        if not issubclass(exc_type, self.expected):
+            # let unexpected exceptions pass through
+            return False
+        # store exception for later retrieval
+        self.exception = exc_value
+        return True
+
+
 class ExpectedException:
     """A context manager to handle expected exceptions.
 

testtools/tests/test_testcase.py

@@ -457,6 +457,45 @@ def test_assertRaisesRegex_wrong_message(self):
             "Observed",
         )
 
+    def test_assertRaises_as_context_manager(self):
+        # assertRaises can be used as a context manager
+        with self.assertRaises(RuntimeError):
+            raise RuntimeError("Error message")
+
+    def test_assertRaises_as_context_manager_no_exception(self):
+        # assertRaises used as context manager fails when no exception is raised
+        with self.assertRaises(self.failureException) as cm:
+            with self.assertRaises(RuntimeError):
+                pass
+        self.assertIn("RuntimeError not raised", str(cm.exception))
+
+    def test_assertRaises_as_context_manager_wrong_exception(self):
+        # assertRaises used as context manager re-raises unexpected exceptions
+        with self.assertRaises(ZeroDivisionError):
+            with self.assertRaises(RuntimeError):
+                raise ZeroDivisionError("Wrong exception")
+
+    def test_assertRaises_as_context_manager_with_msg(self):
+        # assertRaises context manager can accept a msg parameter
+        with self.assertRaises(self.failureException) as cm:
+            with self.assertRaises(RuntimeError, msg="Custom message"):
+                pass
+        self.assertIn("Custom message", str(cm.exception))
+
+    def test_assertRaises_as_context_manager_stores_exception(self):
+        # assertRaises context manager stores the caught exception
+        with self.assertRaises(RuntimeError) as cm:
+            raise RuntimeError("Test error")
+        self.assertIsInstance(cm.exception, RuntimeError)
+        self.assertEqual("Test error", str(cm.exception))
+
+    def test_assertRaises_as_context_manager_with_multiple_exceptions(self):
+        # assertRaises context manager works with multiple exception types
+        exceptions = (RuntimeError, ValueError)
+        with self.assertRaises(exceptions) as cm:
+            raise ValueError("One of the expected exceptions")
+        self.assertIsInstance(cm.exception, ValueError)
+
     def assertFails(self, message, function, *args, **kwargs):
         """Assert that function raises a failure with the given message."""
         failure = self.assertRaises(self.failureException, function, *args, **kwargs)