Merge pull request #170636 from Synthetica9/wait_before_entry

nixos/test-driver: add wait_before_entry

3 files changed, 37 insertions, 8 deletions

diff --git a/nixos/lib/test-driver/test_driver/driver.py b/nixos/lib/test-driver/test_driver/driver.py
index e32f6810ca87f..6542a2e2f6938 100644
--- a/nixos/lib/test-driver/test_driver/driver.py
+++ b/nixos/lib/test-driver/test_driver/driver.py
@@ -220,6 +220,20 @@ class Driver:
                 res = driver.polling_conditions.pop()
                 assert res is self.condition
 
+            def wait(self, timeout: int = 900) -> None:
+                def condition(last: bool) -> bool:
+                    if last:
+                        rootlog.info(f"Last chance for {self.condition.description}")
+                    ret = self.condition.check(force=True)
+                    if not ret and not last:
+                        rootlog.info(
+                            f"({self.condition.description} failure not fatal yet)"
+                        )
+                    return ret
+
+                with rootlog.nested(f"waiting for {self.condition.description}"):
+                    retry(condition, timeout=timeout)
+
         if fun_ is None:
             return Poll
         else:
diff --git a/nixos/lib/test-driver/test_driver/polling_condition.py b/nixos/lib/test-driver/test_driver/polling_condition.py
index 459845452fa12..02ca0a03ab3dc 100644
--- a/nixos/lib/test-driver/test_driver/polling_condition.py
+++ b/nixos/lib/test-driver/test_driver/polling_condition.py
@@ -1,4 +1,5 @@
 from typing import Callable, Optional
+from math import isfinite
 import time
 
 from .logger import rootlog
@@ -14,7 +15,7 @@ class PollingCondition:
     description: Optional[str]
 
     last_called: float
-    entered: bool
+    entry_count: int
 
     def __init__(
         self,
@@ -34,14 +35,21 @@ class PollingCondition:
             self.description = str(description)
 
         self.last_called = float("-inf")
-        self.entered = False
+        self.entry_count = 0
 
-    def check(self) -> bool:
-        if self.entered or not self.overdue:
+    def check(self, force: bool = False) -> bool:
+        if (self.entered or not self.overdue) and not force:
             return True
 
         with self, rootlog.nested(self.nested_message):
-            rootlog.info(f"Time since last: {time.monotonic() - self.last_called:.2f}s")
+            time_since_last = time.monotonic() - self.last_called
+            last_message = (
+                f"Time since last: {time_since_last:.2f}s"
+                if isfinite(time_since_last)
+                else "(not called yet)"
+            )
+
+            rootlog.info(last_message)
             try:
                 res = self.condition()  # type: ignore
             except Exception:
@@ -69,9 +77,16 @@ class PollingCondition:
     def overdue(self) -> bool:
         return self.last_called + self.seconds_interval < time.monotonic()
 
+    @property
+    def entered(self) -> bool:
+        # entry_count should never dip *below* zero
+        assert self.entry_count >= 0
+        return self.entry_count > 0
+
     def __enter__(self) -> None:
-        self.entered = True
+        self.entry_count += 1
 
     def __exit__(self, exc_type, exc_value, traceback) -> None:  # type: ignore
-        self.entered = False
+        assert self.entered
+        self.entry_count -= 1
         self.last_called = time.monotonic()
diff --git a/nixos/tests/vscodium.nix b/nixos/tests/vscodium.nix
index ee884cc4295dd..37bb649889b45 100644
--- a/nixos/tests/vscodium.nix
+++ b/nixos/tests/vscodium.nix
@@ -49,8 +49,8 @@ let
         start_all()
 
         machine.wait_for_unit('graphical.target')
-        machine.wait_until_succeeds('pgrep -x codium')
 
+        codium_running.wait()
         with codium_running:
             # Wait until vscodium is visible. "File" is in the menu bar.
             machine.wait_for_text('Get Started')