1 files changed, 61 insertions, 21 deletions
diff --git a/lib/sqlalchemy/testing/assertions.py b/lib/sqlalchemy/testing/assertions.py
index f5325b0cb..c97202516 100644
--- a/lib/sqlalchemy/testing/assertions.py
+++ b/lib/sqlalchemy/testing/assertions.py
@@ -9,6 +9,7 @@ from __future__ import absolute_import
 
 import contextlib
 import re
+import sys
 import warnings
 
 from . import assertsql
@@ -258,41 +259,80 @@ def eq_ignore_whitespace(a, b, msg=None):
     assert a == b, msg or "%r != %r" % (a, b)
 
 
+def _assert_proper_exception_context(exception):
+    """assert that any exception we're catching does not have a __context__
+    without a __cause__, and that __suppress_context__ is never set.
+
+    Python 3 will report nested as exceptions as "during the handling of
+    error X, error Y occurred". That's not what we want to do.  we want
+    these exceptions in a cause chain.
+
+    """
+
+    if not util.py3k:
+        return
+
+    if (
+        exception.__context__ is not exception.__cause__
+        and not exception.__suppress_context__
+    ):
+        assert False, (
+            "Exception %r was correctly raised but did not set a cause, "
+            "within context %r as its cause."
+            % (exception, exception.__context__)
+        )
+
+
 def assert_raises(except_cls, callable_, *args, **kw):
-    try:
-        callable_(*args, **kw)
-        success = False
-    except except_cls:
-        success = True
+    _assert_raises(except_cls, callable_, args, kw, check_context=True)
 
-    # assert outside the block so it works for AssertionError too !
-    assert success, "Callable did not raise an exception"
+
+def assert_raises_context_ok(except_cls, callable_, *args, **kw):
+    _assert_raises(
+        except_cls, callable_, args, kw,
+    )
 
 
 def assert_raises_return(except_cls, callable_, *args, **kw):
+    return _assert_raises(except_cls, callable_, args, kw, check_context=True)
+
+
+def assert_raises_message(except_cls, msg, callable_, *args, **kwargs):
+    _assert_raises(
+        except_cls, callable_, args, kwargs, msg=msg, check_context=True
+    )
+
+
+def assert_raises_message_context_ok(
+    except_cls, msg, callable_, *args, **kwargs
+):
+    _assert_raises(except_cls, callable_, args, kwargs, msg=msg)
+
+
+def _assert_raises(
+    except_cls, callable_, args, kwargs, msg=None, check_context=False
+):
     ret_err = None
+    if check_context:
+        are_we_already_in_a_traceback = sys.exc_info()[0]
     try:
-        callable_(*args, **kw)
+        callable_(*args, **kwargs)
         success = False
     except except_cls as err:
-        success = True
         ret_err = err
+        success = True
+        if msg is not None:
+            assert re.search(
+                msg, util.text_type(err), re.UNICODE
+            ), "%r !~ %s" % (msg, err,)
+        if check_context and not are_we_already_in_a_traceback:
+            _assert_proper_exception_context(err)
+        print(util.text_type(err).encode("utf-8"))
 
     # assert outside the block so it works for AssertionError too !
     assert success, "Callable did not raise an exception"
-    return ret_err
 
-
-def assert_raises_message(except_cls, msg, callable_, *args, **kwargs):
-    try:
-        callable_(*args, **kwargs)
-        assert False, "Callable did not raise an exception"
-    except except_cls as e:
-        assert re.search(msg, util.text_type(e), re.UNICODE), "%r !~ %s" % (
-            msg,
-            e,
-        )
-        print(util.text_type(e).encode("utf-8"))
+    return ret_err
 
 
 class AssertsCompiledSQL(object):