nixos/test-driver: add `timeout` option for `wait_for_console_text`

Previously, `wait_for_console_text` would block indefinitely until there were lines
shown in the buffer.

This is highly annoying when testing for things that can just hang for some reasons.

This introduces a classical timeout mechanism via non-blocking get on the Queue.

1 files changed, 17 insertions, 4 deletions

diff --git a/nixos/lib/test-driver/test_driver/machine.py b/nixos/lib/test-driver/test_driver/machine.py
index 1a97cedb2e8..3da7b8c96fd 100644
--- a/nixos/lib/test-driver/test_driver/machine.py
+++ b/nixos/lib/test-driver/test_driver/machine.py
@@ -855,17 +855,30 @@ class Machine:
         with self.nested(f"waiting for {regex} to appear on screen"):
             retry(screen_matches)
 
-    def wait_for_console_text(self, regex: str) -> None:
+    def wait_for_console_text(self, regex: str, timeout: float | None = None) -> None:
+        """
+            Wait for the provided regex to appear on console.
+            For each reads,
+
+            If timeout is None, timeout is infinite.
+
+            `timeout` is in seconds.
+        """
         with self.nested(f"waiting for {regex} to appear on console"):
             # Buffer the console output, this is needed
             # to match multiline regexes.
             console = io.StringIO()
+            start = time.time()
             while True:
                 try:
-                    console.write(self.last_lines.get())
+                    # This will return as soon as possible and
+                    # sleep 1 second.
+                    console.write(self.last_lines.get(block=False))
                 except queue.Empty:
-                    self.sleep(1)
-                    continue
+                    time.sleep(1)
+                    if timeout is not None and time.time() - start >= timeout:
+                        # If we reached here, we didn't honor our timeout constraint.
+                        raise Exception(f"`wait_for_console_text` did not match `{regex}` after {timeout} seconds")
                 console.seek(0)
                 matches = re.search(regex, console.read())
                 if matches is not None: