testing: merge code/test_terminal_writer.py into io/test_terminalwriter.py

Author: bluetech

testing/code/test_terminal_writer.py

@@ -1,5 +1,6 @@
 import io
 import os
+import re
 import shutil
 import sys
 from typing import Generator
@@ -207,3 +208,26 @@ def test_combining(self) -> None:
         assert len(text) == 10
         tw.write(text)
         assert tw.width_of_current_line == 9
+
+
+@pytest.mark.parametrize(
+    "has_markup, expected",
+    [
+        pytest.param(
+            True, "{kw}assert{hl-reset} {number}0{hl-reset}\n", id="with markup"
+        ),
+        pytest.param(False, "assert 0\n", id="no markup"),
+    ],
+)
+def test_code_highlight(has_markup, expected, color_mapping):
+    f = io.StringIO()
+    tw = terminalwriter.TerminalWriter(f)
+    tw.hasmarkup = has_markup
+    tw._write_source(["assert 0"])
+    assert f.getvalue().splitlines(keepends=True) == color_mapping.format([expected])
+
+    with pytest.raises(
+        ValueError,
+        match=re.escape("indents size (2) should have same size as lines (1)"),
+    ):
+        tw._write_source(["assert 0"], [" ", " "])