From 9a8b815f2e2b54b31bcc9f5017ce56b1891cacad Mon Sep 17 00:00:00 2001
From: Matt Broadstone <mbroadst@gmail.com>
Date: Fri, 18 Aug 2017 09:14:51 -0400
Subject: [PATCH] feat(mock-server): add the ability to set a message handler

 - modified tests to remove 3s from test run
 - synchronized all destroy calls to remove need for
   timeout-based guess
---
 test/mock/lib/server.js                       | 30 +++++--
 .../functional/rs_mocks/add_remove_tests.js   | 81 ++++++-------------
 2 files changed, 49 insertions(+), 62 deletions(-)

diff --git a/test/mock/lib/server.js b/test/mock/lib/server.js
index fc2ef1ddb..eec2eb4c0 100644
--- a/test/mock/lib/server.js
+++ b/test/mock/lib/server.js
@@ -47,11 +47,21 @@ var Server = function(port, host, options) {
 inherits(Server, EventEmitter);
 
 Server.prototype.destroy = function() {
-  this.state = 'destroyed';
-  this.server.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();
+    });
   });
 };
 
@@ -61,7 +71,7 @@ Server.prototype.start = function() {
   // Return start promise
   return new Promise(function(resolve, reject) {
     self.server.on('error', function(err) {
-      console.log("!!!!!!!!!!!!!!!!!!!! error reject")
+      console.log('!!!!!!!!!!!!!!!!!!!! error reject');
       reject(err);
     });
 
@@ -104,7 +114,11 @@ Server.prototype.start = function() {
 
     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;
diff --git a/test/tests/functional/rs_mocks/add_remove_tests.js b/test/tests/functional/rs_mocks/add_remove_tests.js
index 36dbde886..7318b9f5d 100644
--- a/test/tests/functional/rs_mocks/add_remove_tests.js
+++ b/test/tests/functional/rs_mocks/add_remove_tests.js
@@ -241,7 +241,7 @@ describe('ReplSet Add Remove (mocks)', function() {
     }
   });
 
-  it('Successfully remove a secondary server from the set', {
+  it.only('Successfully remove a secondary server from the set', {
     metadata: {
       requires: {
         generators: true,
@@ -258,7 +258,6 @@ describe('ReplSet Add Remove (mocks)', function() {
       var firstSecondaryServer = null;
       var secondSecondaryServer = null;
       var arbiterServer = null;
-      var running = true;
       var currentIsMasterIndex = 0;
 
       // Default message fields
@@ -365,60 +364,32 @@ describe('ReplSet Add Remove (mocks)', function() {
         secondSecondaryServer = yield mockupdb.createServer(32003, 'localhost');
         arbiterServer = yield mockupdb.createServer(32002, 'localhost');
 
-        // Primary state machine
-        co(function*() {
-          while (running) {
-            var request = yield primaryServer.receive();
-            var doc = request.document;
-
-            if (doc.ismaster) {
-              request.reply(primary[currentIsMasterIndex]);
-            }
+        primaryServer.setMessageHandler(function(request) {
+          var doc = request.document;
+          if (doc.ismaster) {
+            request.reply(primary[currentIsMasterIndex]);
           }
-        }).catch(function() {
-          // console.log(err.stack);
         });
 
-        // First secondary state machine
-        co(function*() {
-          while (running) {
-            var request = yield firstSecondaryServer.receive();
-            var doc = request.document;
-
-            if (doc.ismaster) {
-              request.reply(firstSecondary[currentIsMasterIndex]);
-            }
+        firstSecondaryServer.setMessageHandler(function(request) {
+          var doc = request.document;
+          if (doc.ismaster) {
+            request.reply(firstSecondary[currentIsMasterIndex]);
           }
-        }).catch(function() {
-          // console.log(err.stack);
         });
 
-        // Second secondary state machine
-        co(function*() {
-          while (running) {
-            var request = yield secondSecondaryServer.receive();
-            var doc = request.document;
-
-            if (doc.ismaster) {
-              request.reply(secondSecondary[currentIsMasterIndex]);
-            }
+        secondSecondaryServer.setMessageHandler(function(request) {
+          var doc = request.document;
+          if (doc.ismaster) {
+            request.reply(secondSecondary[currentIsMasterIndex]);
           }
-        }).catch(function() {
-          // console.log(err.stack);
         });
 
-        // Arbiter state machine
-        co(function*() {
-          while (running) {
-            var request = yield arbiterServer.receive();
-            var doc = request.document;
-
-            if (doc.ismaster) {
-              request.reply(arbiter[currentIsMasterIndex]);
-            }
+        arbiterServer.setMessageHandler(function(request) {
+          var doc = request.document;
+          if (doc.ismaster) {
+            request.reply(arbiter[currentIsMasterIndex]);
           }
-        }).catch(function() {
-          // console.log(err.stack);
         });
       });
 
@@ -471,18 +442,16 @@ describe('ReplSet Add Remove (mocks)', function() {
           expect(server.s.replicaSetState.primary).to.not.be.null;
           expect(server.s.replicaSetState.primary.name).to.equal('localhost:32000');
 
-          primaryServer.destroy();
-          firstSecondaryServer.destroy();
-          secondSecondaryServer.destroy();
-          arbiterServer.destroy();
-          server.destroy();
-          running = false;
-
-          setTimeout(function() {
-            expect(Object.keys(Connection.connections())).to.have.length(0);
+          Promise.all([
+            server.destroy(),
+            primaryServer.destroy(),
+            firstSecondaryServer.destroy(),
+            secondSecondaryServer.destroy(),
+            arbiterServer.destroy()
+          ]).then(function() {
             Connection.disableConnectionAccounting();
             done();
-          }, 2000);
+          });
         }
       });