nixos/testing: add interactive serial stdout logs switch and dim them

2 files changed, 28 insertions, 7 deletions

diff --git a/nixos/lib/test-driver/test-driver.py b/nixos/lib/test-driver/test-driver.py
index 270a5969cda5..6b868366e97d 100644
--- a/nixos/lib/test-driver/test-driver.py
+++ b/nixos/lib/test-driver/test-driver.py
@@ -3,6 +3,7 @@ from contextlib import contextmanager, _GeneratorContextManager
 from queue import Queue, Empty
 from typing import Tuple, Any, Callable, Dict, Iterator, Optional, List, Iterable
 from xml.sax.saxutils import XMLGenerator
+from colorama import Style
 import queue
 import io
 import _thread
@@ -151,6 +152,8 @@ class Logger:
         self.xml.startDocument()
         self.xml.startElement("logfile", attrs={})
 
+        self._print_serial_logs = True
+
     def close(self) -> None:
         self.xml.endElement("logfile")
         self.xml.endDocument()
@@ -174,15 +177,21 @@ class Logger:
         self.drain_log_queue()
         self.log_line(message, attributes)
 
-    def enqueue(self, message: Dict[str, str]) -> None:
-        self.queue.put(message)
+    def log_serial(self, message: str, machine: str) -> None:
+        self.enqueue({"msg": message, "machine": machine, "type": "serial"})
+        if self._print_serial_logs:
+            eprint(Style.DIM + "{} # {}".format(machine, message) + Style.RESET_ALL)
+
+    def enqueue(self, item: Dict[str, str]) -> None:
+        self.queue.put(item)
 
     def drain_log_queue(self) -> None:
         try:
             while True:
                 item = self.queue.get_nowait()
-                attributes = {"machine": item["machine"], "type": "serial"}
-                self.log_line(self.sanitise(item["msg"]), attributes)
+                msg = self.sanitise(item["msg"])
+                del item["msg"]
+                self.log_line(msg, item)
         except Empty:
             pass
 
@@ -327,6 +336,9 @@ class Machine:
     def log(self, msg: str) -> None:
         self.logger.log(msg, {"machine": self.name})
 
+    def log_serial(self, msg: str) -> None:
+        self.logger.log_serial(msg, self.name)
+
     def nested(self, msg: str, attrs: Dict[str, str] = {}) -> _GeneratorContextManager:
         my_attrs = {"machine": self.name}
         my_attrs.update(attrs)
@@ -784,8 +796,7 @@ class Machine:
                 # Ignore undecodable bytes that may occur in boot menus
                 line = _line.decode(errors="ignore").replace("\r", "").rstrip()
                 self.last_lines.put(line)
-                eprint("{} # {}".format(self.name, line))
-                self.logger.enqueue({"msg": line, "machine": self.name})
+                self.log_serial(line)
 
         _thread.start_new_thread(process_serial_output, ())
 
@@ -927,6 +938,16 @@ def run_tests() -> None:
             machine.execute("sync")
 
 
+def serial_stdout_on() -> None:
+    global log
+    log._print_serial_logs = True
+
+
+def serial_stdout_off() -> None:
+    global log
+    log._print_serial_logs = False
+
+
 @contextmanager
 def subtest(name: str) -> Iterator[None]:
     with log.nested(name):
diff --git a/nixos/lib/testing-python.nix b/nixos/lib/testing-python.nix
index 7e87313e9639..c9d4f0f0861d 100644
--- a/nixos/lib/testing-python.nix
+++ b/nixos/lib/testing-python.nix
@@ -25,7 +25,7 @@ rec {
       name = "nixos-test-driver";
 
       nativeBuildInputs = [ makeWrapper ];
-      buildInputs = [ (python3.withPackages (p: [ p.ptpython ])) ];
+      buildInputs = [ (python3.withPackages (p: [ p.ptpython p.colorama ])) ];
       checkInputs = with python3Packages; [ pylint black mypy ];
 
       dontUnpack = true;