From 8198617bcce885c280ba5f8acd066c71b311b1ac Mon Sep 17 00:00:00 2001
From: =?utf8?q?=C5=81ukasz=20Langa?= <lukasz@langa.pl>
Date: Wed, 3 Nov 2021 17:47:04 +0100
Subject: [PATCH] [3.9] bpo-45578: add tests for `dis.distb` (GH-29332)
 (#29386)

(cherry picked from commit e346f196819aeb02a8a94205ce3e1536c4c2f105)

Co-authored-by: Nikita Sobolev <mail@sobolevn.me>
---
 Lib/test/test_dis.py                          | 41 +++++++++++++++++++
 .../2021-10-30-19-00-25.bpo-45578.bvu6X2.rst  |  1 +
 2 files changed, 42 insertions(+)
 create mode 100644 Misc/NEWS.d/next/Tests/2021-10-30-19-00-25.bpo-45578.bvu6X2.rst

diff --git a/Lib/test/test_dis.py b/Lib/test/test_dis.py
index ac5836d28897..54bab568ab8a 100644
--- a/Lib/test/test_dis.py
+++ b/Lib/test/test_dis.py
@@ -1202,5 +1202,46 @@ class BytecodeTests(unittest.TestCase):
         b = dis.Bytecode.from_traceback(tb)
         self.assertEqual(b.dis(), dis_traceback)
 
+
+class TestDisTraceback(unittest.TestCase):
+    def setUp(self) -> None:
+        try:  # We need to clean up existing tracebacks
+            del sys.last_traceback
+        except AttributeError:
+            pass
+        return super().setUp()
+
+    def get_disassembly(self, tb):
+        output = io.StringIO()
+        with contextlib.redirect_stdout(output):
+            dis.distb(tb)
+        return output.getvalue()
+
+    def test_distb_empty(self):
+        with self.assertRaises(RuntimeError):
+            dis.distb()
+
+    def test_distb_last_traceback(self):
+        # We need to have an existing last traceback in `sys`:
+        tb = get_tb()
+        sys.last_traceback = tb
+
+        self.assertEqual(self.get_disassembly(None), dis_traceback)
+
+    def test_distb_explicit_arg(self):
+        tb = get_tb()
+
+        self.assertEqual(self.get_disassembly(tb), dis_traceback)
+
+
+class TestDisTracebackWithFile(TestDisTraceback):
+    # Run the `distb` tests again, using the file arg instead of print
+    def get_disassembly(self, tb):
+        output = io.StringIO()
+        with contextlib.redirect_stdout(output):
+            dis.distb(tb, file=output)
+        return output.getvalue()
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/Misc/NEWS.d/next/Tests/2021-10-30-19-00-25.bpo-45578.bvu6X2.rst b/Misc/NEWS.d/next/Tests/2021-10-30-19-00-25.bpo-45578.bvu6X2.rst
new file mode 100644
index 000000000000..3d0e0ca3f04a
--- /dev/null
+++ b/Misc/NEWS.d/next/Tests/2021-10-30-19-00-25.bpo-45578.bvu6X2.rst
@@ -0,0 +1 @@
+Add tests for :func:`dis.distb`
-- 
2.47.3