Improve tests

Author: tomasr8

Lib/test/test_pyrepl/test_pyrepl.py

@@ -35,14 +35,15 @@
     multiline_input,
     code_to_events,
 )
+from _colorize import ANSIColors, get_theme
 from _pyrepl.console import Event
 from _pyrepl.completing_reader import stripcolor
 from _pyrepl._module_completer import (
     ImportParser,
     ModuleCompleter,
     HARDCODED_SUBMODULES,
 )
-from _pyrepl.fancycompleter import Completer as FancyCompleter
+from _pyrepl.fancycompleter import Completer as FancyCompleter, colorize_matches
 import _pyrepl.readline as pyrepl_readline
 from _pyrepl.readline import (
     ReadlineAlikeReader,
@@ -1642,38 +1643,37 @@ def test_suggestions_and_messages(self) -> None:
                         self.assertSetEqual(new_imports, expected_imports)
 
     def test_colorize_import_completions(self) -> None:
-        from _colorize import get_theme
-        from _pyrepl.fancycompleter import colorize_matches
-        from _pyrepl.completing_reader import stripcolor
-
         theme = get_theme()
+        type_color = theme.fancycompleter.type
+        module_color = theme.fancycompleter.module
+        R = ANSIColors.RESET
+
         colorize = lambda names, values: colorize_matches(names, values, theme)
         config = ReadlineConfig(colorize_completions=colorize)
-
         reader = ReadlineAlikeReader(
             console=FakeConsole(events=[]),
             config=config,
         )
 
-        # Multiple completions should be colorized (contain ANSI codes)
-        reader.buffer = list("from collections import d")
+        # "from collections import de" -> defaultdict (type) and deque (type)
+        reader.buffer = list("from collections import de")
         reader.pos = len(reader.buffer)
-        result = reader.get_module_completions()
-        self.assertIsNotNone(result)
-        names, action = result
-        self.assertTrue(len(names) > 1)
-        # Colorized names contain ANSI escape sequences
-        self.assertTrue(any(name != stripcolor(name) for name in names
-                            if name.strip()))
-
-        # Single completion should NOT be colorized
-        reader.buffer = list("from collections import Order")
+        names, action = reader.get_module_completions()
+        self.assertEqual(names, [
+            f"{type_color}defaultdict{R}",
+            f"{type_color}deque{R}",
+        ])
+        self.assertIsNone(action)
+
+        # "from importlib.m" has submodule completions colored as modules
+        reader.buffer = list("from importlib.m")
         reader.pos = len(reader.buffer)
-        result = reader.get_module_completions()
-        self.assertIsNotNone(result)
-        names, action = result
-        self.assertEqual(len(names), 1)
-        self.assertEqual(names[0], stripcolor(names[0]))
+        names, action = reader.get_module_completions()
+        self.assertEqual(names, [
+            f"{module_color}importlib.machinery{R}",
+            f"{module_color}importlib.metadata{R}",
+        ])
+        self.assertIsNone(action)
 
 
 # Audit hook used to check for stdlib modules import side-effects