Bug 1261138 - Allow marionette to have any socket timeout, including None indicating no timeout, r=ato draft
authorJames Graham <james@hoppipolla.co.uk>
Thu, 31 Mar 2016 18:55:14 +0100
changeset 346273 32f9221a4eeb04667fef95b67c7f8917336f6725
parent 346272 cd2480346ac56034766ff54107f82ab237997ec5
child 517396 3f393c3ba1f46a97fe93c4d02e223434a41a3ae2
push id14321
push userbmo:james@hoppipolla.co.uk
push dateThu, 31 Mar 2016 18:00:27 +0000
reviewersato
bugs1261138
milestone48.0a1
Bug 1261138 - Allow marionette to have any socket timeout, including None indicating no timeout, r=ato MozReview-Commit-ID: L6iYzF1cnnX
testing/marionette/client/marionette_driver/marionette.py
testing/marionette/client/marionette_driver/transport.py
--- a/testing/marionette/client/marionette_driver/marionette.py
+++ b/testing/marionette/client/marionette_driver/marionette.py
@@ -535,20 +535,20 @@ class Marionette(object):
     TIMEOUT_PAGE = 'page load'
     DEFAULT_SOCKET_TIMEOUT = 360
     DEFAULT_STARTUP_TIMEOUT = 60
 
     def __init__(self, host='localhost', port=2828, app=None, app_args=None, bin=None,
                  profile=None, addons=None, emulator=None, sdcard=None, emulator_img=None,
                  emulator_binary=None, emulator_res=None, connect_to_running_emulator=False,
                  gecko_log=None, homedir=None, baseurl=None, no_window=False, logdir=None,
-                 busybox=None, symbols_path=None, timeout=None, socket_timeout=None,
-                 device_serial=None, adb_path=None, process_args=None,
-                 adb_host=None, adb_port=None, prefs=None, startup_timeout=None,
-                 workspace=None, verbose=0):
+                 busybox=None, symbols_path=None, timeout=None,
+                 socket_timeout=DEFAULT_SOCKET_TIMEOUT, device_serial=None, adb_path=None,
+                 process_args=None, adb_host=None, adb_port=None, prefs=None,
+                 startup_timeout=None, workspace=None, verbose=0):
         self.host = host
         self.port = self.local_port = port
         self.bin = bin
         self.profile = profile
         self.addons = addons
         self.instance = None
         self.session = None
         self.session_id = None
@@ -556,17 +556,17 @@ class Marionette(object):
         self.chrome_window = None
         self.runner = None
         self.emulator = None
         self.extra_emulators = []
         self.baseurl = baseurl
         self.no_window = no_window
         self._test_name = None
         self.timeout = timeout
-        self.socket_timeout = socket_timeout or self.DEFAULT_SOCKET_TIMEOUT
+        self.socket_timeout = socket_timeout
         self.device_serial = device_serial
         self.adb_host = adb_host
         self.adb_port = adb_port
 
         startup_timeout = startup_timeout or self.DEFAULT_STARTUP_TIMEOUT
 
         if bin:
             port = int(self.port)
--- a/testing/marionette/client/marionette_driver/transport.py
+++ b/testing/marionette/client/marionette_driver/transport.py
@@ -5,16 +5,30 @@
 import datetime
 import errno
 import json
 import socket
 import time
 import types
 
 
+class SocketTimeout(object):
+    def __init__(self, socket, timeout):
+        self.sock = socket
+        self.timeout = timeout
+        self.old_timeout = None
+
+    def __enter__(self):
+        self.old_timeout = self.sock.gettimeout()
+        self.sock.settimeout(self.timeout)
+
+    def __exit__(self, *args, **kwargs):
+        self.sock.settimeout(self.old_timeout)
+
+
 class Message(object):
     def __init__(self, msgid):
         self.id = msgid
 
     def __eq__(self, other):
         return self.id == other.id
 
 
@@ -211,21 +225,20 @@ class TcpTransport(object):
         try:
             self.sock.connect((self.addr, self.port))
         except:
             # Unset self.sock so that the next attempt to send will cause
             # another connection attempt.
             self.sock = None
             raise
 
-        self.sock.settimeout(2.0)
-
-        # first packet is always a JSON Object
-        # which we can use to tell which protocol level we are at
-        raw = self.receive(unmarshal=False)
+        with SocketTimeout(self.sock, 2.0):
+            # first packet is always a JSON Object
+            # which we can use to tell which protocol level we are at
+            raw = self.receive(unmarshal=False)
         hello = json.loads(raw)
         self.protocol = hello.get("marionetteProtocol", 1)
         self.application_type = hello.get("applicationType")
 
         return (self.protocol, self.application_type)
 
     def send(self, obj):
         """Send message to the remote server.  Allowed input is a