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.

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])