nixos/test-driver: integrate Python unittest assertions

Replaces / Closes #345948

I tried to integrate `pytest` assertions because I like the reporting,
but I only managed to get the very basic thing and even that was messing
around a lot with its internals.

The approach in #345948 shifts too much maintenance effort to us, so
it's not really desirable either.

After discussing with Benoit on Ocean Sprint about this, we decided that
it's probably the best compromise to integrate `unittest`: it also
provides good diffs when needed, but the downside is that existing tests
don't benefit from it.

This patch essentially does the following things:

* Add a new global `t` that is an instance of a `unittest.TestCase`
  class. I decided to just go for `t` given that e.g.
  `tester.assertEqual` (or any other longer name) seems quite verbose.

* Use a special class for errors that get special treatment:
  * The traceback is minimized to only include frames from the
    testScript: in this case I don't really care about anything else and
    IMHO that's just visual noise.

    This is not the case for other exceptions since these may indicate a
    bug and then people should be able to send the full traceback to the
    maintainers.
  * Display the error, but with `!!!` as prefix to make sure it's
    easier to spot in between other logs.

This looks e.g. like

    !!! Traceback (most recent call last):
    !!!   File "<string>", line 7, in <module>
    !!!     foo()
    !!!   File "<string>", line 5, in foo
    !!!     t.assertEqual({"foo":[1,2,{"foo":"bar"}]},{"foo":[1,2,{"bar":"foo"}],"bar":[1,2,3,4,"foo"]})
    !!!
    !!! NixOSAssertionError: {'foo': [1, 2, {'foo': 'bar'}]} != {'foo': [1, 2, {'bar': 'foo'}], 'bar': [1, 2, 3, 4, 'foo']}
    !!! - {'foo': [1, 2, {'foo': 'bar'}]}
    !!! + {'bar': [1, 2, 3, 4, 'foo'], 'foo': [1, 2, {'bar': 'foo'}]}
    cleanup
    kill machine (pid 9)
    qemu-system-x86_64: terminating on signal 15 from pid 6 (/nix/store/wz0j2zi02rvnjiz37nn28h3gfdq61svz-python3-3.12.9/bin/python3.12)
    kill vlan (pid 7)
    (finished: cleanup, in 0.00 seconds)

Co-authored-by: bew <bew@users.noreply.github.com>

nixos/doc/manual/development/writing-nixos-tests.section.md

@@ -121,8 +121,7 @@ and checks that the output is more-or-less correct:
 ```py
 machine.start()
 machine.wait_for_unit("default.target")
-if not "Linux" in machine.succeed("uname"):
-  raise Exception("Wrong OS")
+t.assertIn("Linux", machine.succeed("uname"), "Wrong OS")
 ```
 
 The first line is technically unnecessary; machines are implicitly started
@@ -134,6 +133,8 @@ starting them in parallel:
 start_all()
 ```
 
+Under the variable `t`, all assertions from [`unittest.TestCase`](https://docs.python.org/3/library/unittest.html) are available.
+
 If the hostname of a node contains characters that can't be used in a
 Python variable name, those characters will be replaced with
 underscores in the variable name, so `nodes.machine-a` will be exposed

nixos/lib/test-driver/src/test_driver/driver.py

@@ -1,12 +1,15 @@
 import os
 import re
 import signal
+import sys
 import tempfile
 import threading
+import traceback
 from collections.abc import Callable, Iterator
 from contextlib import AbstractContextManager, contextmanager
 from pathlib import Path
 from typing import Any
+from unittest import TestCase
 
 from test_driver.logger import AbstractLogger
 from test_driver.machine import Machine, NixStartScript, retry
@@ -38,6 +41,14 @@ def pythonize_name(name: str) -> str:
     return re.sub(r"^[^A-z_]|[^A-z0-9_]", "_", name)
 
 
+class NixOSAssertionError(AssertionError):
+    pass
+
+
+class Tester(TestCase):
+    failureException = NixOSAssertionError
+
+
 class Driver:
     """A handle to the driver that sets up the environment
     and runs the tests"""
@@ -140,6 +151,7 @@ class Driver:
             serial_stdout_on=self.serial_stdout_on,
             polling_condition=self.polling_condition,
             Machine=Machine,  # for typing
+            t=Tester(),
         )
         machine_symbols = {pythonize_name(m.name): m for m in self.machines}
         # If there's exactly one machine, make it available under the name
@@ -163,7 +175,31 @@ class Driver:
         """Run the test script"""
         with self.logger.nested("run the VM test script"):
             symbols = self.test_symbols()  # call eagerly
-            exec(self.tests, symbols, None)
+            try:
+                exec(self.tests, symbols, None)
+            except NixOSAssertionError:
+                exc_type, exc, tb = sys.exc_info()
+                filtered = [
+                    frame
+                    for frame in traceback.extract_tb(tb)
+                    if frame.filename == "<string>"
+                ]
+
+                self.logger.log_test_error("Traceback (most recent call last):")
+                code = self.tests.splitlines()
+                for frame, line in zip(filtered, traceback.format_list(filtered)):
+                    self.logger.log_test_error(line.rstrip())
+                    if lineno := frame.lineno:
+                        self.logger.log_test_error(
+                            f"    {code[lineno - 1].strip()}",
+                        )
+
+                self.logger.log_test_error("") # blank line for readability
+                exc_prefix = exc_type.__name__ if exc_type is not None else "Error"
+                for line in f"{exc_prefix}: {exc}".splitlines():
+                    self.logger.log_test_error(line)
+
+                sys.exit(1)
 
     def run_tests(self) -> None:
         """Run the test script (for non-interactive test runs)"""

nixos/lib/test-driver/src/test_driver/logger.py

@@ -44,6 +44,10 @@ class AbstractLogger(ABC):
     def error(self, *args, **kwargs) -> None:  # type: ignore
         pass
 
+    @abstractmethod
+    def log_test_error(self, *args, **kwargs) -> None:  # type:ignore
+        pass
+
     @abstractmethod
     def log_serial(self, message: str, machine: str) -> None:
         pass
@@ -97,6 +101,9 @@ class JunitXMLLogger(AbstractLogger):
         self.tests[self.currentSubtest].stderr += args[0] + os.linesep
         self.tests[self.currentSubtest].failure = True
 
+    def log_test_error(self, *args, **kwargs) -> None:  # type: ignore
+        self.error(*args, **kwargs)
+
     def log_serial(self, message: str, machine: str) -> None:
         if not self._print_serial_logs:
             return
@@ -156,6 +163,10 @@ class CompositeLogger(AbstractLogger):
         for logger in self.logger_list:
             logger.warning(*args, **kwargs)
 
+    def log_test_error(self, *args, **kwargs) -> None:  # type: ignore
+        for logger in self.logger_list:
+            logger.log_test_error(*args, **kwargs)
+
     def error(self, *args, **kwargs) -> None:  # type: ignore
         for logger in self.logger_list:
             logger.error(*args, **kwargs)
@@ -222,6 +233,11 @@ class TerminalLogger(AbstractLogger):
 
         self._eprint(Style.DIM + f"{machine} # {message}" + Style.RESET_ALL)
 
+    def log_test_error(self, *args, **kwargs) -> None:  # type: ignore
+        prefix = Fore.RED + "!!! " + Style.RESET_ALL
+        # NOTE: using `warning` instead of `error` to ensure it does not exit after printing the first log
+        self.warning(f"{prefix}{args[0]}", *args[1:], **kwargs)
+
 
 class XMLLogger(AbstractLogger):
     def __init__(self, outfile: str) -> None:
@@ -261,6 +277,9 @@ class XMLLogger(AbstractLogger):
     def error(self, *args, **kwargs) -> None:  # type: ignore
         self.log(*args, **kwargs)
 
+    def log_test_error(self, *args, **kwargs) -> None:  # type: ignore
+        self.log(*args, **kwargs)
+
     def log(self, message: str, attributes: dict[str, str] = {}) -> None:
         self.drain_log_queue()
         self.log_line(message, attributes)

nixos/lib/test-script-prepend.py

@@ -8,6 +8,7 @@ from test_driver.logger import AbstractLogger
 from typing import Callable, Iterator, ContextManager, Optional, List, Dict, Any, Union
 from typing_extensions import Protocol
 from pathlib import Path
+from unittest import TestCase
 
 
 class RetryProtocol(Protocol):
@@ -51,3 +52,4 @@ join_all: Callable[[], None]
 serial_stdout_off: Callable[[], None]
 serial_stdout_on: Callable[[], None]
 polling_condition: PollingConditionProtocol
+t: TestCase