Unittest - Allow specifying custom timeout for slow commands.

* Useful when testing with big tasks (1000+ annotations) or sync
operations that take longer than 1 second (default)

test/basetest/task.py

@@ -177,7 +177,8 @@ class Task(object):
         with open(self.taskrc, 'w') as rc:
             rc.write("data.location={0}\n"
                      "confirmation=no\n"
-                     "hooks=off\n".format(self.datadir))
+                     "hooks=off\n"
+                     "".format(self.datadir))
 
         # Setup configuration to talk to taskd automatically
         if self.taskd is not None:
@@ -255,7 +256,8 @@ class Task(object):
         cmd = (self.taskw, "config", "--", var, value)
         return run_cmd_wait(cmd, env=self.env)
 
-    def runSuccess(self, args=(), input=None, merge_streams=True):
+    def runSuccess(self, args=(), input=None, merge_streams=True,
+                   timeout=1):
         """Invoke task with given arguments and fail if exit code != 0
 
         Use runError if you want exit_code to be tested automatically and
@@ -267,6 +269,9 @@ class Task(object):
 
         If merge_streams=True stdout and stderr will be merged into stdout.
 
+        timeout = number of seconds the test will wait for every task call.
+        Defaults to 1 second if not specified. Unit is seconds.
+
         Returns (exit_code, stdout, stderr)
         """
         # Create a copy of the command
@@ -274,14 +279,16 @@ class Task(object):
         command.extend(args)
 
         output = run_cmd_wait_nofail(command, input,
-                                     merge_streams=merge_streams, env=self.env)
+                                     merge_streams=merge_streams,
+                                     env=self.env,
+                                     timeout=timeout)
 
         if output[0] != 0:
             raise CommandError(command, *output)
 
         return output
 
-    def runError(self, args=(), input=None, merge_streams=True):
+    def runError(self, args=(), input=None, merge_streams=True, timeout=1):
         """Invoke task with given arguments and fail if exit code == 0
 
         Use runSuccess if you want exit_code to be tested automatically and
@@ -293,6 +300,9 @@ class Task(object):
 
         If merge_streams=True stdout and stderr will be merged into stdout.
 
+        timeout = number of seconds the test will wait for every task call.
+        Defaults to 1 second if not specified. Unit is seconds.
+
         Returns (exit_code, stdout, stderr)
         """
         # Create a copy of the command
@@ -300,7 +310,9 @@ class Task(object):
         command.extend(args)
 
         output = run_cmd_wait_nofail(command, input,
-                                     merge_streams=merge_streams, env=self.env)
+                                     merge_streams=merge_streams,
+                                     env=self.env,
+                                     timeout=timeout)
 
         # output[0] is the exit code
         if output[0] == 0 or output[0] is None:

test/basetest/utils.py

@@ -46,6 +46,9 @@ def binary_location(cmd):
 def wait_process(proc, timeout=1):
     """Wait for process to finish
     """
+    if timeout is None:
+        timeout = 1
+
     sleeptime = .1
     # Max number of attempts until giving up
     tries = int(timeout / sleeptime)
@@ -62,7 +65,7 @@ def wait_process(proc, timeout=1):
     return exit
 
 
-def _get_output(proc, input):
+def _get_output(proc, input, timeout=None):
     """Collect output from the subprocess without blocking the main process if
     subprocess hangs.
     """
@@ -84,7 +87,7 @@ def _get_output(proc, input):
     t.start()
 
     # A task process shouldn't take longer than 1 second to finish
-    exit = wait_process(proc)
+    exit = wait_process(proc, timeout)
 
     # If it does take longer than 1 second, abort it
     if exit is None:
@@ -113,7 +116,7 @@ def _get_output(proc, input):
 
 
 def run_cmd_wait(cmd, input=None, stdout=PIPE, stderr=PIPE,
-                 merge_streams=False, env=os.environ):
+                 merge_streams=False, env=os.environ, timeout=None):
     "Run a subprocess and wait for it to finish"
 
     if input is None:
@@ -128,7 +131,7 @@ def run_cmd_wait(cmd, input=None, stdout=PIPE, stderr=PIPE,
 
     p = Popen(cmd, stdin=stdin, stdout=stdout, stderr=stderr, bufsize=1,
               close_fds=ON_POSIX, env=env)
-    out, err, exit = _get_output(p, input)
+    out, err, exit = _get_output(p, input, timeout)
 
     if exit != 0:
         raise CommandError(cmd, exit, out, err)
@@ -276,6 +279,7 @@ except ImportError:
                         return name
         return None
 
+
 def parse_datafile(file):
     """Parse .data files on the client and server treating files as JSON
     """