nixos/test-driver: run callbacks during Machine.wait_for_console_text

.github/CODEOWNERS

@@ -75,7 +75,7 @@
 /nixos/modules/system                                 @dasJ
 
 # NixOS integration test driver
-/nixos/lib/test-driver  @tfc
+/nixos/lib/test-driver  @tfc @synthetica9
 
 # Updaters
 ## update.nix

nixos/lib/test-driver/default.nix

@@ -14,7 +14,7 @@
 
 python3Packages.buildPythonApplication rec {
   pname = "nixos-test-driver";
-  version = "1.1";
+  version = "1.2";
   src = ./.;
 
   propagatedBuildInputs = [ coreutils netpbm python3Packages.colorama python3Packages.ptpython qemu_pkg socat vde2 ]
@@ -29,4 +29,9 @@ python3Packages.buildPythonApplication rec {
     pylint --errors-only --enable=unused-import ${src}/test_driver
     black --check --diff ${src}/test_driver
   '';
+
+  meta = {
+    maintainers = with lib.maintainers; [ synthetica tfc ];
+    homepage = "https://nixos.org/manual/nixos/unstable/index.html#sec-nixos-tests";
+  };
 }

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

@@ -2,6 +2,7 @@
 from pathlib import Path
 from typing import Any, Dict, Iterator, List, Union, Optional, Callable, ContextManager
 import os
+import re
 import tempfile
 
 from test_driver.logger import rootlog
@@ -10,6 +11,29 @@
 from test_driver.polling_condition import PollingCondition
 
 
+@contextmanager
+def must_raise(
+    regex: str,
+    exception: type[BaseException] = Exception,
+) -> Iterator[None]:
+    short_desc = f"{exception.__name__} {regex!r}"
+    msg = f"Expected {short_desc}"
+
+    with rootlog.nested(f"Waiting for {short_desc}"):
+        try:
+            yield
+        except exception as e:
+            e_str = str(e)
+            if re.search(regex, str(e)):
+                return
+            raise Exception(f"{msg}, but string {str(e)!r} did not match") from e
+        except Exception as e:
+            raise Exception(
+                f"{msg}, but {e!r} is not an instance of {exception!r}"
+            ) from e
+        raise Exception(f"{msg}, but no exception was raised")
+
+
 class Driver:
     """A handle to the driver that sets up the environment
     and runs the tests"""
@@ -91,6 +115,7 @@ def subtest(name: str) -> Iterator[None]:
             serial_stdout_on=self.serial_stdout_on,
             polling_condition=self.polling_condition,
             Machine=Machine,  # for typing
+            must_raise=must_raise,
         )
         machine_symbols = {m.name: m for m in self.machines}
         # If there's exactly one machine, make it available under the name

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

@@ -805,21 +805,22 @@ def screen_matches(last: bool) -> bool:
         with self.nested("waiting for {} to appear on screen".format(regex)):
             retry(screen_matches)
 
-    def wait_for_console_text(self, regex: str) -> None:
+    def wait_for_console_text(self, regex: str) -> re.Match[str]:
         with self.nested("waiting for {} to appear on console".format(regex)):
             # Buffer the console output, this is needed
             # to match multiline regexes.
             console = io.StringIO()
             while True:
                 try:
-                    console.write(self.last_lines.get())
+                    console.write(self.last_lines.get(timeout=1))
                 except queue.Empty:
+                    self.run_callbacks()
                     self.sleep(1)
                     continue
                 console.seek(0)
                 matches = re.search(regex, console.read())
                 if matches is not None:
-                    return
+                    return matches
 
     def send_key(self, key: str) -> None:
         key = CHAR_TO_KEY.get(key, key)