1 files changed, 34 insertions, 2 deletions
diff --git a/Lib/test/test_code_module.py b/Lib/test/test_code_module.py
index 7a80a80..3394b39 100644
--- a/Lib/test/test_code_module.py
+++ b/Lib/test/test_code_module.py
@@ -1,6 +1,7 @@
 "Test InteractiveConsole and InteractiveInterpreter from code module"
 import sys
 import unittest
+from textwrap import dedent
 from contextlib import ExitStack
 from unittest import mock
 from test import support
@@ -78,9 +79,40 @@ class TestInteractiveConsole(unittest.TestCase):
         self.console.interact(banner='')
         self.assertEqual(len(self.stderr.method_calls), 1)
 
+    def test_cause_tb(self):
+        self.infunc.side_effect = ["raise ValueError('') from AttributeError",
+                                    EOFError('Finished')]
+        self.console.interact()
+        output = ''.join(''.join(call[1]) for call in self.stderr.method_calls)
+        expected = dedent("""
+        AttributeError
+
+        The above exception was the direct cause of the following exception:
+
+        Traceback (most recent call last):
+          File "<console>", line 1, in <module>
+        ValueError
+        """)
+        self.assertIn(expected, output)
+
+    def test_context_tb(self):
+        self.infunc.side_effect = ["try: ham\nexcept: eggs\n",
+                                    EOFError('Finished')]
+        self.console.interact()
+        output = ''.join(''.join(call[1]) for call in self.stderr.method_calls)
+        expected = dedent("""
+        Traceback (most recent call last):
+          File "<console>", line 1, in <module>
+        NameError: name 'ham' is not defined
+
+        During handling of the above exception, another exception occurred:
+
+        Traceback (most recent call last):
+          File "<console>", line 2, in <module>
+        NameError: name 'eggs' is not defined
+        """)
+        self.assertIn(expected, output)
 
-def test_main():
-    support.run_unittest(TestInteractiveConsole)
 
 if __name__ == "__main__":
     unittest.main()