refactor(mock-server): improve test run time across the board

.travis.yml

@@ -1,19 +1,34 @@
 sudo: required
 dist: trusty
-language: node_js
 
-node_js:
-  - "4"
-  - "6"
-  - "8"
+language: node_js
 
 cache:
   yarn: true
   directories:
     - node_modules
 
-env:
-  - MONGODB_VERSION=2.6.x
-  - MONGODB_VERSION=3.0.x
-  - MONGODB_VERSION=3.2.x
-  - MONGODB_VERSION=3.4.x
+jobs:
+  include:
+    - stage: tests
+      node_js: 4
+      env: MONGODB_VERSION=2.6.x
+    - stage: tests
+      node_js: 4
+      env: MONGODB_VERSION=3.0.x
+    - stage: tests
+      node_js: 4
+      env: MONGODB_VERSION=3.2.x
+    - stage: tests
+      node_js: 4
+      env: MONGODB_VERSION=3.4.x
+    - stage: tests
+      node_js: 4
+      env: MONGODB_VERSION=3.5.x
+
+    - stage: tests
+      node_js: 6
+      env: MONGODB_VERSION=3.4.x
+    - stage: tests
+      node_js: 8
+      env: MONGODB_VERSION=3.4.x

lib/connection/connection.js

@@ -36,6 +36,8 @@ var debugFields = [
   'promoteBuffers',
   'checkServerIdentity'
 ];
+
+var connectionAccountingSpy = undefined;
 var connectionAccounting = false;
 var connections = {};
 
@@ -185,13 +187,18 @@ Connection.prototype.resetSocketTimeout = function() {
   }
 };
 
-Connection.enableConnectionAccounting = function() {
+Connection.enableConnectionAccounting = function(spy) {
+  if (spy) {
+    connectionAccountingSpy = spy;
+  }
+
   connectionAccounting = true;
   connections = {};
 };
 
 Connection.disableConnectionAccounting = function() {
   connectionAccounting = false;
+  connectionAccountingSpy = undefined;
 };
 
 Connection.connections = function() {
@@ -201,11 +208,19 @@ Connection.connections = function() {
 function deleteConnection(id) {
   // console.log("=== deleted connection " + id + " :: " + (connections[id] ? connections[id].port : ''))
   delete connections[id];
+
+  if (connectionAccountingSpy) {
+    connectionAccountingSpy.deleteConnection(id);
+  }
 }
 
 function addConnection(id, connection) {
   // console.log("=== added connection " + id + " :: " + connection.port)
   connections[id] = connection;
+
+  if (connectionAccountingSpy) {
+    connectionAccountingSpy.addConnection(id, connection);
+  }
 }
 
 //

test/mock/index.js

@@ -1,11 +1,47 @@
 var Server = require('./lib/server');
 
+const cleanup = (servers, spy, callback) => {
+  if (typeof spy === 'function') {
+    callback = spy;
+    spy = undefined;
+  }
+
+  if (!Array.isArray(servers)) {
+    throw new Error('First argument must be an array of mock servers');
+  }
+
+  if (spy) {
+    const alreadyDrained = spy.connectionCount() === 0;
+    const finish = () => {
+      callback(null, null);
+    };
+
+    if (!alreadyDrained) {
+      spy.once('drained', () => finish());
+    }
+
+    const cleanupPromise = Promise.all(servers.map(server => server.destroy())).catch(err =>
+      callback(err, null)
+    );
+
+    if (alreadyDrained) {
+      cleanupPromise.then(() => finish());
+    }
+  } else {
+    Promise.all(servers.map(server => server.destroy()))
+      .then(() => callback(null, null))
+      .catch(err => callback(err, null));
+  }
+};
+
 /*
  * Main module
  */
 module.exports = {
   createServer: function(port, host, options) {
     options = options || {};
     return new Server(port, host, options).start();
-  }
+  },
+
+  cleanup: cleanup
 };

test/mock/lib/server.js

@@ -33,7 +33,7 @@ var Server = function(port, host, options) {
   this.host = host;
   this.port = port;
   // Create a server socket
-  this.socket = net.createServer();
+  this.server = net.createServer();
   // Responses
   this.messages = [];
   // state
@@ -47,11 +47,21 @@ var Server = function(port, host, options) {
 inherits(Server, EventEmitter);
 
 Server.prototype.destroy = function() {
-  this.state = 'destroyed';
-  this.socket.close();
+  var self = this;
+  if (self.state === 'destroyed') {
+    return Promise.resolve();
+  }
+
+  return new Promise(function(resolve, reject) {
+    self.sockets.forEach(function(x) {
+      x.destroy();
+    });
 
-  this.sockets.forEach(function(x) {
-    x.destroy();
+    self.server.close(function(err) {
+      if (err) return reject(err);
+      self.state = 'destroyed';
+      resolve();
+    });
   });
 };
 
@@ -60,12 +70,12 @@ Server.prototype.start = function() {
 
   // Return start promise
   return new Promise(function(resolve, reject) {
-    self.socket.on('error', function(err) {
+    self.server.on('error', function(err) {
       console.log('!!!!!!!!!!!!!!!!!!!! error reject');
       reject(err);
     });
 
-    self.socket.on('connection', function(c) {
+    self.server.on('connection', function(c) {
       self.connections = self.connections + 1;
       self.sockets.push(c);
 
@@ -98,13 +108,17 @@ Server.prototype.start = function() {
       });
     });
 
-    self.socket.listen(self.port, self.host, function() {
+    self.server.listen(self.port, self.host, function() {
       resolve(self);
     });
 
     self.on('message', function(message, connection) {
       var request = new Request(self, connection, message);
-      self.messages.push(request);
+      if (self.messageHandler) {
+        self.messageHandler(request);
+      } else {
+        self.messages.push(request);
+      }
     });
 
     self.state = 'running';
@@ -133,6 +147,10 @@ Server.prototype.receive = function() {
   });
 };
 
+Server.prototype.setMessageHandler = function(messageHandler) {
+  this.messageHandler = messageHandler;
+};
+
 var protocol = function(self, message) {
   var index = 0;
   self.isCompressed = false;