Refactor tests

Author: ambv

Lib/test/test_pyrepl/support.py

@@ -6,14 +6,24 @@
 from _pyrepl.console import Console, Event
 from _pyrepl.readline import ReadlineAlikeReader, ReadlineConfig
 from _pyrepl.simple_interact import _strip_final_indent
+from _pyrepl.utils import unbracket, ANSI_ESCAPE_SEQUENCE
+
+
+class ScreenEqualMixin:
+    def assert_screen_equal(
+        self, reader: ReadlineAlikeReader, expected: str, clean: bool = False
+    ):
+        actual = clean_screen(reader) if clean else reader.screen
+        expected = expected.split("\n")
+        self.assertListEqual(actual, expected)
 
 
 def multiline_input(reader: ReadlineAlikeReader, namespace: dict | None = None):
     saved = reader.more_lines
     try:
         reader.more_lines = partial(more_lines, namespace=namespace)
-        reader.ps1 = reader.ps2 = ">>>"
-        reader.ps3 = reader.ps4 = "..."
+        reader.ps1 = reader.ps2 = ">>> "
+        reader.ps3 = reader.ps4 = "... "
         return reader.readline()
     finally:
         reader.more_lines = saved
@@ -38,18 +48,22 @@ def code_to_events(code: str):
         yield Event(evt="key", data=c, raw=bytearray(c.encode("utf-8")))
 
 
-def clean_screen(screen: Iterable[str]):
+def clean_screen(reader: ReadlineAlikeReader) -> list[str]:
     """Cleans color and console characters out of a screen output.
 
     This is useful for screen testing, it increases the test readability since
     it strips out all the unreadable side of the screen.
     """
     output = []
-    for line in screen:
-        if line.startswith(">>>") or line.startswith("..."):
-            line = line[3:]
+    for line in reader.screen:
+        line = unbracket(line, including_content=True)
+        line = ANSI_ESCAPE_SEQUENCE.sub("", line)
+        for prefix in (reader.ps1, reader.ps2, reader.ps3, reader.ps4):
+            if line.startswith(prefix):
+                line = line[len(prefix):]
+                break
         output.append(line)
-    return "\n".join(output).strip()
+    return output
 
 
 def prepare_reader(console: Console, **kwargs):
@@ -99,6 +113,9 @@ def handle_all_events(
     prepare_console=partial(prepare_console, width=10),
 )
 
+reader_no_colors = partial(prepare_reader, can_colorize=False)
+reader_force_colors = partial(prepare_reader, can_colorize=True)
+
 
 class FakeConsole(Console):
     def __init__(self, events, encoding="utf-8") -> None:

Lib/test/test_pyrepl/test_pyrepl.py

@@ -17,12 +17,12 @@
 
 from .support import (
     FakeConsole,
+    ScreenEqualMixin,
     handle_all_events,
     handle_events_narrow_console,
     more_lines,
     multiline_input,
     code_to_events,
-    clean_screen,
 )
 from _pyrepl.console import Event
 from _pyrepl.readline import (ReadlineAlikeReader, ReadlineConfig,
@@ -587,7 +587,7 @@ def test_auto_indent_ignore_comments(self):
         self.assertEqual(output, output_code)
 
 
-class TestPyReplOutput(TestCase):
+class TestPyReplOutput(ScreenEqualMixin, TestCase):
     def prepare_reader(self, events):
         console = FakeConsole(events)
         config = ReadlineConfig(readline_completer=None)
@@ -620,7 +620,7 @@ def test_basic(self):
 
         output = multiline_input(reader)
         self.assertEqual(output, "1+1")
-        self.assertEqual(clean_screen(reader.screen), "1+1")
+        self.assert_screen_equal(reader, "1+1", clean=True)
 
     def test_get_line_buffer_returns_str(self):
         reader = self.prepare_reader(code_to_events("\n"))
@@ -654,11 +654,13 @@ def test_multiline_edit(self):
         reader = self.prepare_reader(events)
 
         output = multiline_input(reader)
-        self.assertEqual(output, "def f():\n    ...\n    ")
-        self.assertEqual(clean_screen(reader.screen), "def f():\n    ...")
+        expected = "def f():\n    ...\n    "
+        self.assertEqual(output, expected)
+        self.assert_screen_equal(reader, expected, clean=True)
         output = multiline_input(reader)
-        self.assertEqual(output, "def g():\n    pass\n    ")
-        self.assertEqual(clean_screen(reader.screen), "def g():\n    pass")
+        expected = "def g():\n    pass\n    "
+        self.assertEqual(output, expected)
+        self.assert_screen_equal(reader, expected, clean=True)
 
     def test_history_navigation_with_up_arrow(self):
         events = itertools.chain(
@@ -677,16 +679,16 @@ def test_history_navigation_with_up_arrow(self):
 
         output = multiline_input(reader)
         self.assertEqual(output, "1+1")
-        self.assertEqual(clean_screen(reader.screen), "1+1")
+        self.assert_screen_equal(reader, "1+1", clean=True)
         output = multiline_input(reader)
         self.assertEqual(output, "2+2")
-        self.assertEqual(clean_screen(reader.screen), "2+2")
+        self.assert_screen_equal(reader, "2+2", clean=True)
         output = multiline_input(reader)
         self.assertEqual(output, "2+2")
-        self.assertEqual(clean_screen(reader.screen), "2+2")
+        self.assert_screen_equal(reader, "2+2", clean=True)
         output = multiline_input(reader)
         self.assertEqual(output, "1+1")
-        self.assertEqual(clean_screen(reader.screen), "1+1")
+        self.assert_screen_equal(reader, "1+1", clean=True)
 
     def test_history_with_multiline_entries(self):
         code = "def foo():\nx = 1\ny = 2\nz = 3\n\ndef bar():\nreturn 42\n\n"
@@ -705,11 +707,9 @@ def test_history_with_multiline_entries(self):
         output = multiline_input(reader)
         output = multiline_input(reader)
         output = multiline_input(reader)
-        self.assertEqual(
-            clean_screen(reader.screen),
-            'def foo():\n    x = 1\n    y = 2\n    z = 3'
-        )
-        self.assertEqual(output, "def foo():\n    x = 1\n    y = 2\n    z = 3\n    ")
+        expected = "def foo():\n    x = 1\n    y = 2\n    z = 3\n    "
+        self.assert_screen_equal(reader, expected, clean=True)
+        self.assertEqual(output, expected)
 
 
     def test_history_navigation_with_down_arrow(self):
@@ -728,7 +728,7 @@ def test_history_navigation_with_down_arrow(self):
 
         output = multiline_input(reader)
         self.assertEqual(output, "1+1")
-        self.assertEqual(clean_screen(reader.screen), "1+1")
+        self.assert_screen_equal(reader, "1+1", clean=True)
 
     def test_history_search(self):
         events = itertools.chain(
@@ -745,23 +745,23 @@ def test_history_search(self):
 
         output = multiline_input(reader)
         self.assertEqual(output, "1+1")
-        self.assertEqual(clean_screen(reader.screen), "1+1")
+        self.assert_screen_equal(reader, "1+1", clean=True)
         output = multiline_input(reader)
         self.assertEqual(output, "2+2")
-        self.assertEqual(clean_screen(reader.screen), "2+2")
+        self.assert_screen_equal(reader, "2+2", clean=True)
         output = multiline_input(reader)
         self.assertEqual(output, "3+3")
-        self.assertEqual(clean_screen(reader.screen), "3+3")
+        self.assert_screen_equal(reader, "3+3", clean=True)
         output = multiline_input(reader)
         self.assertEqual(output, "1+1")
-        self.assertEqual(clean_screen(reader.screen), "1+1")
+        self.assert_screen_equal(reader, "1+1", clean=True)
 
     def test_control_character(self):
         events = code_to_events("c\x1d\n")
         reader = self.prepare_reader(events)
         output = multiline_input(reader)
         self.assertEqual(output, "c\x1d")
-        self.assertEqual(clean_screen(reader.screen), "c")
+        self.assert_screen_equal(reader, "c\x1d", clean=True)
 
     def test_history_search_backward(self):
         # Test <page up> history search backward with "imp" input
@@ -781,7 +781,7 @@ def test_history_search_backward(self):
         # search for "imp" in history
         output = multiline_input(reader)
         self.assertEqual(output, "import os")
-        self.assertEqual(clean_screen(reader.screen), "import os")
+        self.assert_screen_equal(reader, "import os", clean=True)
 
     def test_history_search_backward_empty(self):
         # Test <page up> history search backward with an empty input
@@ -800,7 +800,7 @@ def test_history_search_backward_empty(self):
         # search backward in history
         output = multiline_input(reader)
         self.assertEqual(output, "import os")
-        self.assertEqual(clean_screen(reader.screen), "import os")
+        self.assert_screen_equal(reader, "import os", clean=True)
 
 
 class TestPyReplCompleter(TestCase):