bpo-24959: fix unittest.assertRaises bug where traceback entries are dropped from chained exceptions (GH-23688)

2 files changed, 94 insertions, 14 deletions

diff --git a/Lib/unittest/result.py b/Lib/unittest/result.py
index ce7468e..3da7005 100644
--- a/Lib/unittest/result.py
+++ b/Lib/unittest/result.py
@@ -173,18 +173,10 @@ class TestResult(object):
     def _exc_info_to_string(self, err, test):
         """Converts a sys.exc_info()-style tuple of values into a string."""
         exctype, value, tb = err
-        # Skip test runner traceback levels
-        while tb and self._is_relevant_tb_level(tb):
-            tb = tb.tb_next
-
-        if exctype is test.failureException:
-            # Skip assert*() traceback levels
-            length = self._count_relevant_tb_levels(tb)
-        else:
-            length = None
+        tb = self._clean_tracebacks(exctype, value, tb, test)
         tb_e = traceback.TracebackException(
             exctype, value, tb,
-            limit=length, capture_locals=self.tb_locals, compact=True)
+            capture_locals=self.tb_locals, compact=True)
         msgLines = list(tb_e.format())
 
         if self.buffer:
@@ -200,16 +192,49 @@ class TestResult(object):
                 msgLines.append(STDERR_LINE % error)
         return ''.join(msgLines)
 
+    def _clean_tracebacks(self, exctype, value, tb, test):
+        ret = None
+        first = True
+        excs = [(exctype, value, tb)]
+        while excs:
+            (exctype, value, tb) = excs.pop()
+            # Skip test runner traceback levels
+            while tb and self._is_relevant_tb_level(tb):
+                tb = tb.tb_next
+
+            # Skip assert*() traceback levels
+            if exctype is test.failureException:
+                self._remove_unittest_tb_frames(tb)
+
+            if first:
+                ret = tb
+                first = False
+            else:
+                value.__traceback__ = tb
+
+            if value is not None:
+                for c in (value.__cause__, value.__context__):
+                    if c is not None:
+                        excs.append((type(c), c, c.__traceback__))
+        return ret
 
     def _is_relevant_tb_level(self, tb):
         return '__unittest' in tb.tb_frame.f_globals
 
-    def _count_relevant_tb_levels(self, tb):
-        length = 0
+    def _remove_unittest_tb_frames(self, tb):
+        '''Truncates usercode tb at the first unittest frame.
+
+        If the first frame of the traceback is in user code,
+        the prefix up to the first unittest frame is returned.
+        If the first frame is already in the unittest module,
+        the traceback is not modified.
+        '''
+        prev = None
         while tb and not self._is_relevant_tb_level(tb):
-            length += 1
+            prev = tb
             tb = tb.tb_next
-        return length
+        if prev is not None:
+            prev.tb_next = None
 
     def __repr__(self):
         return ("<%s run=%i errors=%i failures=%i>" %
diff --git a/Lib/unittest/test/test_result.py b/Lib/unittest/test/test_result.py
index 224a784..c616f28 100644
--- a/Lib/unittest/test/test_result.py
+++ b/Lib/unittest/test/test_result.py
@@ -220,6 +220,61 @@ class Test_TestResult(unittest.TestCase):
         self.assertIs(test_case, test)
         self.assertIsInstance(formatted_exc, str)
 
+    def test_addFailure_filter_traceback_frames(self):
+        class Foo(unittest.TestCase):
+            def test_1(self):
+                pass
+
+        test = Foo('test_1')
+        def get_exc_info():
+            try:
+                test.fail("foo")
+            except:
+                return sys.exc_info()
+
+        exc_info_tuple = get_exc_info()
+
+        full_exc = traceback.format_exception(*exc_info_tuple)
+
+        result = unittest.TestResult()
+        result.startTest(test)
+        result.addFailure(test, exc_info_tuple)
+        result.stopTest(test)
+
+        formatted_exc = result.failures[0][1]
+        dropped = [l for l in full_exc if l not in formatted_exc]
+        self.assertEqual(len(dropped), 1)
+        self.assertIn("raise self.failureException(msg)", dropped[0])
+
+    def test_addFailure_filter_traceback_frames_context(self):
+        class Foo(unittest.TestCase):
+            def test_1(self):
+                pass
+
+        test = Foo('test_1')
+        def get_exc_info():
+            try:
+                try:
+                    test.fail("foo")
+                except:
+                    raise ValueError(42)
+            except:
+                return sys.exc_info()
+
+        exc_info_tuple = get_exc_info()
+
+        full_exc = traceback.format_exception(*exc_info_tuple)
+
+        result = unittest.TestResult()
+        result.startTest(test)
+        result.addFailure(test, exc_info_tuple)
+        result.stopTest(test)
+
+        formatted_exc = result.failures[0][1]
+        dropped = [l for l in full_exc if l not in formatted_exc]
+        self.assertEqual(len(dropped), 1)
+        self.assertIn("raise self.failureException(msg)", dropped[0])
+
     # "addError(test, err)"
     # ...
     # "Called when the test case test raises an unexpected exception err