diff --git a/mozdevice/mozdevice/devicemanager.py b/mozdevice/mozdevice/devicemanager.py index 9c5ce3e..25cd410 100644 --- a/mozdevice/mozdevice/devicemanager.py +++ b/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: diff --git a/mozdevice/mozdevice/devicemanagerADB.py b/mozdevice/mozdevice/devicemanagerADB.py index 779c771..950b74f 100644 --- a/mozdevice/mozdevice/devicemanagerADB.py +++ b/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: # 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: diff --git a/mozdevice/mozdevice/devicemanagerSUT.py b/mozdevice/mozdevice/devicemanagerSUT.py index d1e7ae4..31ff59e 100644 --- a/mozdevice/mozdevice/devicemanagerSUT.py +++ b/mozdevice/mozdevice/devicemanagerSUT.py @@ -265,16 +265,16 @@ def _doCmds(self, cmdlist, outputfile, timeout): # returns: # success: # 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