Merge pull request #234513 from NixOS/test-driver/wait_for_console_timeout

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

1 files changed, 29 insertions, 14 deletions

diff --git a/nixos/lib/test-driver/test_driver/machine.py b/nixos/lib/test-driver/test_driver/machine.py
index 20cfde8af4b..ec36c05cbbb 100644
--- a/nixos/lib/test-driver/test_driver/machine.py
+++ b/nixos/lib/test-driver/test_driver/machine.py
@@ -856,21 +856,36 @@ 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: int | 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.
+        """
+        # Buffer the console output, this is needed
+        # to match multiline regexes.
+        console = io.StringIO()
+        def console_matches() -> bool:
+            nonlocal console
+            try:
+                # This will return as soon as possible and
+                # sleep 1 second.
+                console.write(self.last_lines.get(block=False))
+            except queue.Empty:
+                pass
+            console.seek(0)
+            matches = re.search(regex, console.read())
+            return (matches is not None)
+
         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()
-            while True:
-                try:
-                    console.write(self.last_lines.get())
-                except queue.Empty:
-                    self.sleep(1)
-                    continue
-                console.seek(0)
-                matches = re.search(regex, console.read())
-                if matches is not None:
-                    return
+            if timeout is not None:
+                retry(console_matches, timeout)
+            else:
+                while not console_matches():
+                    pass
 
     def send_key(
         self, key: str, delay: Optional[float] = 0.01, log: Optional[bool] = True