add timeout to ADB.cmd()

(cherry picked from commit 39039f6)

airtest/core/android/adb.py

@@ -9,21 +9,33 @@
 import subprocess
 import threading
 
-from six import PY3, text_type, binary_type
+from six import PY3, text_type, binary_type, raise_from
 from six.moves import reduce
 
 from airtest.core.android.constant import (DEFAULT_ADB_PATH, IP_PATTERN,
                                            SDK_VERISON_ANDROID7)
 from airtest.core.error import (AdbError, AdbShellError, AirtestError,
-                                DeviceConnectionError)
-from airtest.utils.compat import decode_path
+                                DeviceConnectionError, AdbTimeoutExpired)
+from airtest.utils.compat import decode_path, raisefrom
 from airtest.utils.logger import get_logger
 from airtest.utils.nbsp import NonBlockingStreamReader
 from airtest.utils.retry import retries
 from airtest.utils.snippet import get_std_encoding, reg_cleanup, split_cmd
 
 LOGGING = get_logger(__name__)
 
+if sys.platform.startswith("win"):
+    # Don't display the Windows GPF dialog if the invoked program dies.
+    try:
+        SUBPROCESS_FLAG = subprocess.CREATE_NO_WINDOW  # in Python 3.7+
+    except AttributeError:
+        import ctypes
+        SEM_NOGPFAULTERRORBOX = 0x0002  # From MSDN
+        ctypes.windll.kernel32.SetErrorMode(SEM_NOGPFAULTERRORBOX)  # win32con.CREATE_NO_WINDOW?
+        SUBPROCESS_FLAG = 0x8000000
+else:
+    SUBPROCESS_FLAG = 0
+
 
 class ADB(object):
     """adb client object class"""
@@ -150,18 +162,20 @@ def start_cmd(self, cmds, device=True):
             cmds,
             stdin=subprocess.PIPE,
             stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE
+            stderr=subprocess.PIPE,
+            creationflags=SUBPROCESS_FLAG
         )
         return proc
 
-    def cmd(self, cmds, device=True, ensure_unicode=True):
+    def cmd(self, cmds, device=True, ensure_unicode=True, timeout=None):
         """
         Run the adb command(s) in subprocess and return the standard output
 
         Args:
             cmds: command(s) to be run
             device: if True, the device serial number must be specified by -s serialno argument
             ensure_unicode: encode/decode unicode of standard outputs (stdout, stderr)
+            timeout: timeout in seconds
 
         Raises:
             DeviceConnectionError: if any error occurs when connecting the device
@@ -172,7 +186,26 @@ def cmd(self, cmds, device=True, ensure_unicode=True):
 
         """
         proc = self.start_cmd(cmds, device)
-        stdout, stderr = proc.communicate()
+        if timeout:
+            if sys.version_info[:2] >= (3, 3):
+                # in Python 3.3+
+                try:
+                    stdout, stderr = proc.communicate(timeout=timeout)
+                except subprocess.TimeoutExpired as e:
+                    proc.kill()
+                    stdout, stderr = proc.communicate()
+                    raise_from(AdbTimeoutExpired(stdout, stderr, e), None)
+            else:
+                timer = threading.Timer(timeout, proc.kill)
+                try:
+                    timer.start()
+                    stdout, stderr = proc.communicate()
+                finally:
+                    timer.cancel()
+                    if proc.returncode > 0:
+                        raise AdbTimeoutExpired(stdout, stderr, ("Command '%s' timed out after %s seconds" % (cmds, timeout)))
+        else:
+            stdout, stderr = proc.communicate()
 
         if ensure_unicode:
             stdout = stdout.decode(get_std_encoding(sys.stdout))
@@ -289,14 +322,10 @@ def wait_for_device(self, timeout=5):
             None
 
         """
-        proc = self.start_cmd("wait-for-device")
-        timer = threading.Timer(timeout, proc.kill)
-        timer.start()
-        ret = proc.wait()
-        if ret == 0:
-            timer.cancel()
-        else:
-            raise DeviceConnectionError("device not ready")
+        try:
+            self.cmd("wait-for-device", timeout=timeout)
+        except AdbTimeoutExpired as e:
+            raisefrom(DeviceConnectionError, "device not ready", e)
 
     def start_shell(self, cmds):
         """

airtest/core/error.py

@@ -63,6 +63,22 @@ class AdbShellError(AdbError):
     pass
 
 
+class AdbTimeoutExpired(Exception):
+    """
+    This exception is raised when the timeout expires while waiting for ADB command.
+    """
+    def __init__(self, stdout, stderr, exp=None):
+        self.stdout = stdout
+        self.stderr = stderr
+        self.exp = exp
+
+    def __str__(self):
+        if self.exp:
+            return self.exp.__str__() + "\n stdout[%s] stderr[%s]" % (self.stdout, self.stderr)
+        else:
+            return "stdout[%s] stderr[%s]" % (self.stdout, self.stderr)
+
+
 class DeviceConnectionError(BaseError):
     """
         device connection error

airtest/utils/compat.py

@@ -1,6 +1,6 @@
 import sys
 import os
-from six import PY3
+from six import PY3, raise_from, reraise
 
 
 EXT = ".air"  # script dir extension
@@ -33,3 +33,10 @@ def script_log_dir(script_path, logdir):
     elif logdir:
         logdir = decode_path(logdir)
     return logdir
+
+
+def raisefrom(exc_type, message, exc):
+    if sys.version_info[:2] >= (3, 2):
+        raise_from(exc_type(message), exc)
+    else:
+        reraise(exc_type, '%s - %s' % (message, exc), sys.exc_info()[2])