Bug 1182681 - Raise IOError when receiving empty string on socket, r=jgraham

* * *
Bug 1182681 - Raise IOError when we receive an emtry string, r=jgraham

--HG--
extra : commitid : 2JF1C6eDB0l

testing/marionette/driver/marionette_driver/marionette.py

@@ -631,8 +631,7 @@ class Marionette(object):
         self.client = MarionetteTransport(
             self.host,
             self.port,
-            self.socket_timeout,
-            instance=self.instance)
+            self.socket_timeout)
 
         if emulator:
             if busybox:
@@ -690,6 +689,12 @@ class Marionette(object):
 
         try:
             response = self.client.send(message)
+        except IOError:
+            if self.instance and not hasattr(self.instance, 'detached'):
+                # If we've launched the binary we've connected to, wait
+                # for it to shut down.
+                returncode = self.instance.runner.wait()
+                raise IOError("process died with returncode %d" % returncode)
         except socket.timeout:
             self.session = None
             self.window = None

testing/marionette/transport/marionette_transport/transport.py

@@ -20,7 +20,7 @@ class MarionetteTransport(object):
     max_packet_length = 4096
     connection_lost_msg = "Connection to Marionette server is lost. Check gecko.log (desktop firefox) or logcat (b2g) for errors."
 
-    def __init__(self, addr, port, socket_timeout=360.0, instance=None):
+    def __init__(self, addr, port, socket_timeout=360.0):
         self.addr = addr
         self.port = port
         self.socket_timeout = socket_timeout
@@ -28,7 +28,6 @@ class MarionetteTransport(object):
         self.traits = None
         self.applicationType = None
         self.actor = 'root'
-        self.instance = instance
 
     def _recv_n_bytes(self, n):
         """ Convenience method for receiving exactly n bytes from
@@ -53,16 +52,13 @@ class MarionetteTransport(object):
         bytes_to_recv = 10
         while time.time() - now < self.socket_timeout:
             try:
-                response += self.sock.recv(bytes_to_recv)
+                data = self.sock.recv(bytes_to_recv)
+                response += data
             except socket.timeout:
                 pass
-            if self.instance and not hasattr(self.instance, 'detached'):
-                # If we've launched the binary we've connected to, make
-                # sure it hasn't died.
-                poll = self.instance.runner.process_handler.proc.poll()
-                if poll is not None:
-                    # process isn't alive
-                    raise IOError("process has died with return code %d" % poll)
+            else:
+                if not data:
+                    raise IOError(self.connection_lost_msg)
             sep = response.find(':')
             if sep > -1:
                 length = response[0:sep]
@@ -70,7 +66,7 @@ class MarionetteTransport(object):
                 if len(remaining) == int(length):
                     return json.loads(remaining)
                 bytes_to_recv = int(length) - len(remaining)
-        raise IOError(self.connection_lost_msg)
+        raise socket.timeout('connection timed out after %d s' % self.socket_timeout)
 
     def connect(self):
         """ Connect to the server and process the hello message we expect