diff --git a/nixos/lib/test-driver/test-driver.py b/nixos/lib/test-driver/test-driver.py index 270a5969cda..6b868366e97 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 7e87313e963..c9d4f0f0861 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;