Merge pull request #115 from hfiguiere/bug1033402

Bug 1033402 - waitForSocket() before connecting to marionette. r=jlal

lib/marionette/drivers/tcp-sync.js

@@ -34,31 +34,84 @@ TcpSync.prototype.setScriptTimeout = function(timeout) {
   this.sockit.setPollTimeout(timeout + SOCKET_TIMEOUT_EXTRA);
 };
 
+/**
+ * Utility to wait for the marionette socket to be ready.
+ *
+ * @method waitForSocket
+ * @param {Object} [options] for timeout.
+ * @param {Number} [options.interval] time between running test.
+ * @param {Number} [options.timeout] maximum wallclock time before
+ *   failing test.
+ * @param {Function} [callback] callback.
+ */
+TcpSync.prototype.waitForSocket = function(options, callback) {
+  // the logic is lifted from http://dxr.mozilla.org/mozilla-central/source/testing/marionette/client/marionette/marionette.py#562
+
+  if (typeof(options) === 'function') {
+    callback = options;
+    options = null;
+  }
+
+  options = options || {};
+  var interval = options.interval || 1000;
+  var timeout = options.timeout || 10000;
+
+  var sockit = new sockittome.Sockit();
+  var start = Date.now();
+  var self = this;
+
+  function probeSocket() {
+    try {
+      sockit.connect({ host: this.host, port: this.port });
+
+      var s = sockit.read(16).toString();
+      sockit.close();
+      if (s.indexOf(":") != -1) {
+        callback();
+        return;
+      }
+    }
+    catch(e) {
+    }
+    // timeout. Abort.
+    if ((Date.now() - start) > timeout) {
+      console.error('timeout connecting to b2g.');
+      return;
+    }
+    // interval delay for the next iteration.
+    setTimeout(probeSocket.bind(self), interval);
+  };
+
+  probeSocket();
+}
+
 TcpSync.prototype.connect = function(callback) {
 
-  try {
-    this.sockit.connect({ host: this.host, port: this.port });
-  } catch(err) {
-    if (Date.now() - this._beginConnect >= this.connectionTimeout) {
-      callback(err);
+  this.waitForSocket(function _connect() {
+    try {
+      this.sockit.connect({ host: this.host, port: this.port });
+    } catch(err) {
+      if (Date.now() - this._beginConnect >= this.connectionTimeout) {
+        callback(err);
+      }
+      setTimeout(_connect.bind(this, callback), this.retryInterval);
+      return;
     }
-    setTimeout(this.connect.bind(this, callback), this.retryInterval);
-    return;
-  }
 
-  if (!this._beginConnect) {
-    this._beginConnect = Date.now();
-  }
+    if (!this._beginConnect) {
+      this._beginConnect = Date.now();
+    }
 
-  // Ensure this method's resolution is asynchronous in all cases
-  setTimeout(function() {
-    debug('socket connected');
-    delete this._beginConnect;
+    // Ensure this method's resolution is asynchronous in all cases
+    process.nextTick(function() {
+      debug('socket connected');
+      delete this._beginConnect;
 
-    this._readResponse();
+      this._readResponse();
 
-    callback();
-  }.bind(this), 0);
+      callback();
+    }.bind(this));
+  }.bind(this));
 };
 
 TcpSync.prototype.defaultCallback = function(err, result) {

test/marionette/drivers/tcp-sync-test-server.js

@@ -0,0 +1,32 @@
+var net = require('net');
+var debug = require('debug')('marionette:tcp-sync-test-server');
+
+const response = '53:{"from":"root","applicationType":"gecko","traits":[]}';
+
+var server = net.createServer(function(connection) { //'connection' listener
+  debug('server connected');
+  connection.on('end', function() {
+    debug('server disconnected');
+  });
+  connection.on('error', function(e) {
+    debug('error', e);
+  });
+  connection.write(response);
+  connection.pipe(connection);
+});
+
+
+function runServer(port) {
+  server.listen(port, function() { //'listening' listener
+    debug('server bound');
+  });
+}
+
+
+if (require.main === module) {
+  try {
+    runServer(1234);
+  }
+  catch(e) {
+  }
+}

test/marionette/drivers/tcp-sync-test.js

@@ -0,0 +1,31 @@
+suite('marionette/drivers/tcp-sync', function() {
+
+  var child = null;
+
+  teardown(function() {
+    if (child) {
+      child.kill();
+      child = null;
+    }
+  });
+
+  test('can wait for marionette socket', function(done) {
+
+    this.timeout(5000);
+
+    var TcpSync = require('../../../lib/marionette/drivers/tcp-sync.js');
+    tcpSync = new TcpSync({port: 1234, host: 'localhost'});
+
+    tcpSync.waitForSocket( {timeout: 3000} , function() {
+      done();
+    });
+
+    setTimeout(function () {
+      child = require('child_process')
+        .fork('./test/marionette/drivers/tcp-sync-test-server.js');
+    }, 1000);
+
+  });
+
+});
+