nixos/test-driver: add junit-xml logger

We add a new logger that allows generating a junit-xml compatible report
listing the subtests used in the nixos integration test. Junit-xml is a
widely used standard for test reports. The report can be used for quick
evaluation of which subtest failed.

2 files changed, 90 insertions, 3 deletions

diff --git a/nixos/lib/test-driver/test_driver/__init__.py b/nixos/lib/test-driver/test_driver/__init__.py
index 9daae1e941a65..91da765e0d2b1 100755
--- a/nixos/lib/test-driver/test_driver/__init__.py
+++ b/nixos/lib/test-driver/test_driver/__init__.py
@@ -6,7 +6,7 @@ from pathlib import Path
 import ptpython.repl
 
 from test_driver.driver import Driver
-from test_driver.logger import rootlog
+from test_driver.logger import JunitXMLLogger, XMLLogger, rootlog
 
 
 class EnvDefault(argparse.Action):
@@ -93,6 +93,11 @@ def main() -> None:
         type=writeable_dir,
     )
     arg_parser.add_argument(
+        "--junit-xml",
+        help="Enable JunitXML report generation to the given path",
+        type=Path,
+    )
+    arg_parser.add_argument(
         "testscript",
         action=EnvDefault,
         envvar="testScript",
@@ -102,6 +107,14 @@ def main() -> None:
 
     args = arg_parser.parse_args()
 
+    output_directory = args.output_directory.resolve()
+
+    if "LOGFILE" in os.environ.keys():
+        rootlog.add_logger(XMLLogger(os.environ["LOGFILE"]))
+
+    if args.junit_xml:
+        rootlog.add_logger(JunitXMLLogger(output_directory / args.junit_xml))
+
     if not args.keep_vm_state:
         rootlog.info("Machine state will be reset. To keep it, pass --keep-vm-state")
 
@@ -109,7 +122,7 @@ def main() -> None:
         args.start_scripts,
         args.vlans,
         args.testscript.read_text(),
-        args.output_directory.resolve(),
+        output_directory,
         args.keep_vm_state,
         args.global_timeout,
     ) as driver:
diff --git a/nixos/lib/test-driver/test_driver/logger.py b/nixos/lib/test-driver/test_driver/logger.py
index f50e3b78a8280..ec1e25bf2db42 100644
--- a/nixos/lib/test-driver/test_driver/logger.py
+++ b/nixos/lib/test-driver/test_driver/logger.py
@@ -1,3 +1,4 @@
+import atexit
 import codecs
 import os
 import sys
@@ -5,12 +6,14 @@ import time
 import unicodedata
 from abc import ABC, abstractmethod
 from contextlib import ExitStack, contextmanager
+from pathlib import Path
 from queue import Empty, Queue
 from typing import Any, Dict, Iterator, List
 from xml.sax.saxutils import XMLGenerator
 from xml.sax.xmlreader import AttributesImpl
 
 from colorama import Fore, Style
+from junit_xml import TestCase, TestSuite
 
 
 class AbstractLogger(ABC):
@@ -49,6 +52,77 @@ class AbstractLogger(ABC):
         pass
 
 
+class JunitXMLLogger(AbstractLogger):
+
+    class TestCaseState:
+        def __init__(self) -> None:
+            self.stdout = ""
+            self.stderr = ""
+            self.failure = False
+
+    def __init__(self, outfile: Path) -> None:
+        self.tests: dict[str, JunitXMLLogger.TestCaseState] = {
+            "main": self.TestCaseState()
+        }
+        self.currentSubtest = "main"
+        self.outfile: Path = outfile
+        self._print_serial_logs = True
+        atexit.register(self.close)
+
+    def log(self, message: str, attributes: Dict[str, str] = {}) -> None:
+        self.tests[self.currentSubtest].stdout += message + os.linesep
+
+    @contextmanager
+    def subtest(self, name: str, attributes: Dict[str, str] = {}) -> Iterator[None]:
+        old_test = self.currentSubtest
+        self.tests.setdefault(name, self.TestCaseState())
+        self.currentSubtest = name
+
+        yield
+
+        self.currentSubtest = old_test
+
+    @contextmanager
+    def nested(self, message: str, attributes: Dict[str, str] = {}) -> Iterator[None]:
+        self.log(message)
+        yield
+
+    def info(self, *args, **kwargs) -> None:  # type: ignore
+        self.tests[self.currentSubtest].stdout += args[0] + os.linesep
+
+    def warning(self, *args, **kwargs) -> None:  # type: ignore
+        self.tests[self.currentSubtest].stdout += args[0] + os.linesep
+
+    def error(self, *args, **kwargs) -> None:  # type: ignore
+        self.tests[self.currentSubtest].stderr += args[0] + os.linesep
+        self.tests[self.currentSubtest].failure = True
+
+    def log_serial(self, message: str, machine: str) -> None:
+        if not self._print_serial_logs:
+            return
+
+        self.log(f"{machine} # {message}")
+
+    def print_serial_logs(self, enable: bool) -> None:
+        self._print_serial_logs = enable
+
+    def close(self) -> None:
+        with open(self.outfile, "w") as f:
+            test_cases = []
+            for name, test_case_state in self.tests.items():
+                tc = TestCase(
+                    name,
+                    stdout=test_case_state.stdout,
+                    stderr=test_case_state.stderr,
+                )
+                if test_case_state.failure:
+                    tc.add_failure_info("test case failed")
+
+                test_cases.append(tc)
+            ts = TestSuite("NixOS integration test", test_cases)
+            f.write(TestSuite.to_xml_string([ts]))
+
+
 class CompositeLogger(AbstractLogger):
     def __init__(self, logger_list: List[AbstractLogger]) -> None:
         self.logger_list = logger_list
@@ -238,4 +312,4 @@ class XMLLogger(AbstractLogger):
 
 terminal_logger = TerminalLogger()
 xml_logger = XMLLogger()
-rootlog: AbstractLogger = CompositeLogger([terminal_logger, xml_logger])
+rootlog: CompositeLogger = CompositeLogger([terminal_logger, xml_logger])