nixos/test-driver: add `timeout` option for `wait_for_console_text` (variant 2)

1 files changed, 21 insertions, 19 deletions

diff --git a/nixos/lib/test-driver/test_driver/machine.py b/nixos/lib/test-driver/test_driver/machine.py
index 3da7b8c96fdcd..3673271798a8f 100644
--- a/nixos/lib/test-driver/test_driver/machine.py
+++ b/nixos/lib/test-driver/test_driver/machine.py
@@ -855,7 +855,7 @@ class Machine:
         with self.nested(f"waiting for {regex} to appear on screen"):
             retry(screen_matches)
 
-    def wait_for_console_text(self, regex: str, timeout: float | None = None) -> 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,
@@ -864,25 +864,27 @@ class Machine:
 
             `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()
-            start = time.time()
-            while True:
-                try:
-                    # This will return as soon as possible and
-                    # sleep 1 second.
-                    console.write(self.last_lines.get(block=False))
-                except queue.Empty:
-                    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:
-                    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