nixos/test-driver: add wait_before_entry to polling_condition

1 files changed, 8 insertions, 0 deletions

diff --git a/nixos/lib/test-driver/test_driver/driver.py b/nixos/lib/test-driver/test_driver/driver.py
index e32f6810ca8..19b88975988 100644
--- a/nixos/lib/test-driver/test_driver/driver.py
+++ b/nixos/lib/test-driver/test_driver/driver.py
@@ -202,6 +202,7 @@ class Driver:
         *,
         seconds_interval: float = 2.0,
         description: Optional[str] = None,
+        wait_before_entry: bool = False,
     ) -> Union[Callable[[Callable], ContextManager], ContextManager]:
         driver = self
 
@@ -212,8 +213,15 @@ class Driver:
                     seconds_interval,
                     description,
                 )
+                self.wait_before_entry = wait_before_entry
 
             def __enter__(self) -> None:
+                if self.wait_before_entry:
+                    with rootlog.nested(
+                        f"waiting before entering polling condition {self.condition.description}"
+                    ):
+                        retry(lambda x: self.condition.check(force=True))
+
                 driver.polling_conditions.append(self.condition)
 
             def __exit__(self, a, b, c) -> None:  # type: ignore