From e62b8020f3d6597ffe4c5444fe824546af88e739 Mon Sep 17 00:00:00 2001
From: Cole Helbling <cole.helbling@determinate.systems>
Date: Mon, 15 Nov 2021 12:23:25 -0800
Subject: [PATCH] nixos/test-driver: add (functional) timeouts to more
 functions

A retry timeout doesn't really help if the thing it's retrying may block
forever.
---
 nixos/lib/test-driver/test-driver.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/nixos/lib/test-driver/test-driver.py b/nixos/lib/test-driver/test-driver.py
index 1d45e001bb4c..b6fbc03c1777 100755
--- a/nixos/lib/test-driver/test-driver.py
+++ b/nixos/lib/test-driver/test-driver.py
@@ -634,12 +634,12 @@ class Machine:
             pass_fds=[self.shell.fileno()],
         )
 
-    def succeed(self, *commands: str) -> str:
+    def succeed(self, *commands: str, timeout: Optional[int] = None) -> str:
         """Execute each command and check that it succeeds."""
         output = ""
         for command in commands:
             with self.nested("must succeed: {}".format(command)):
-                (status, out) = self.execute(command)
+                (status, out) = self.execute(command, timeout=timeout)
                 if status != 0:
                     self.log("output: {}".format(out))
                     raise Exception(
@@ -648,12 +648,12 @@ class Machine:
                 output += out
         return output
 
-    def fail(self, *commands: str) -> str:
+    def fail(self, *commands: str, timeout: Optional[int] = None) -> str:
         """Execute each command and check that it fails."""
         output = ""
         for command in commands:
             with self.nested("must fail: {}".format(command)):
-                (status, out) = self.execute(command)
+                (status, out) = self.execute(command, timeout=timeout)
                 if status == 0:
                     raise Exception(
                         "command `{}` unexpectedly succeeded".format(command)
@@ -669,14 +669,14 @@ class Machine:
 
         def check_success(_: Any) -> bool:
             nonlocal output
-            status, output = self.execute(command)
+            status, output = self.execute(command, timeout=timeout)
             return status == 0
 
         with self.nested("waiting for success: {}".format(command)):
             retry(check_success, timeout)
             return output
 
-    def wait_until_fails(self, command: str) -> str:
+    def wait_until_fails(self, command: str, timeout: int = 900) -> str:
         """Wait until a command returns failure.
         Throws an exception on timeout.
         """
@@ -684,7 +684,7 @@ class Machine:
 
         def check_failure(_: Any) -> bool:
             nonlocal output
-            status, output = self.execute(command)
+            status, output = self.execute(command, timeout=timeout)
             return status != 0
 
         with self.nested("waiting for failure: {}".format(command)):