Add timeouts for shell/checkCmd

mozdevice/mozdevice/devicemanager.py

@@ -39,7 +39,7 @@ def not_implemented(*args, **kwargs):
 class DeviceManager:
 
   @abstractmethod
-  def shell(self, cmd, outputfile, env=None, cwd=None):
+  def shell(self, cmd, outputfile, env=None, cwd=None, timeout=None):
     """
     executes shell command on device
     returns:

mozdevice/mozdevice/devicemanagerADB.py

@@ -8,6 +8,7 @@
 import os
 import sys
 import tempfile
+import time
 
 class DeviceManagerADB(DeviceManager):
 
@@ -94,7 +95,7 @@ def __del__(self):
   # returns:
   # success: <return code>
   # failure: None
-  def shell(self, cmd, outputfile, env=None, cwd=None):
+  def shell(self, cmd, outputfile, env=None, cwd=None, timeout=None):
     # FIXME: this function buffers all output of the command into memory,
     # always. :(
 
@@ -117,6 +118,15 @@ def shell(self, cmd, outputfile, env=None, cwd=None):
     args.extend(["shell", cmdline])
     proc = subprocess.Popen(args,
                             stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    if timeout:
+        start_time = time.time()
+        ret_code = proc.poll()
+        while ((time.time() - start_time) <= timeout) and ret_code == None:
+            time.sleep(1)
+            ret_code = proc.poll()
+        if ret_code == None:
+            proc.kill()
+            raise DMError("Timeout exceeded for shell call")
     (stdout, stderr) = proc.communicate()
     outputfile.write(stdout.rstrip('\n'))
 
@@ -710,7 +720,7 @@ def runCmdAs(self, args):
       args.insert(2, self.packageName)
     return self.runCmd(args)
 
-  def checkCmd(self, args):
+  def checkCmd(self, args, timeout=None):
     # If we are not root but have run-as, and we're trying to execute 
     # a shell command then using run-as is the best we can do
     finalArgs = [self.adbPath]
@@ -720,13 +730,24 @@ def checkCmd(self, args):
       args.insert(1, "run-as")
       args.insert(2, self.packageName)
     finalArgs.extend(args)
+    if timeout:
+        proc = subprocess.Popen(finalArgs)
+        start_time = time.time()
+        ret_code = proc.poll()
+        while ((time.time() - start_time) <= timeout) and ret_code == None:
+            time.sleep(1)
+            ret_code = proc.poll()
+        if ret_code == None:
+            proc.kill()
+            raise DMError("Timeout exceeded for checkCmd call")
+        return ret_code
     return subprocess.check_call(finalArgs)
 
-  def checkCmdAs(self, args):
+  def checkCmdAs(self, args, timeout=None):
     if (self.useRunAs):
       args.insert(1, "run-as")
       args.insert(2, self.packageName)
-    return self.checkCmd(args)
+    return self.checkCmd(args, timeout)
 
   # external function
   # returns:

mozdevice/mozdevice/devicemanagerSUT.py

@@ -265,16 +265,16 @@ def _doCmds(self, cmdlist, outputfile, timeout):
   # returns:
   # success: <return code>
   # failure: None
-  def shell(self, cmd, outputfile, env=None, cwd=None):
+  def shell(self, cmd, outputfile, env=None, cwd=None, timeout=None):
     cmdline = self._escapedCommandLine(cmd)
     if env:
       cmdline = '%s %s' % (self.formatEnvString(env), cmdline)
 
     try:
       if cwd:
-        self.sendCmds([{ 'cmd': 'execcwd %s %s' % (cwd, cmdline) }], outputfile)
+        self.sendCmds([{ 'cmd': 'execcwd %s %s' % (cwd, cmdline) }], outputfile, timeout)
       else:
-        self.sendCmds([{ 'cmd': 'exec su -c "%s"' % cmdline }], outputfile)
+        self.sendCmds([{ 'cmd': 'exec su -c "%s"' % cmdline }], outputfile, timeout)
     except AgentError:
       return None