From ab3b70b9fb982840551cb5d9450458adb1f74f8f Mon Sep 17 00:00:00 2001
From: Matt Broadstone <mbroadst@gmail.com>
Date: Fri, 15 Sep 2017 14:54:11 -0400
Subject: [PATCH] feat(mock): support a means of consistently cleaning up mock
 servers

NODE-1132
---
 test/mock/index.js | 33 ++++++++++++++++++++++++++++++++-
 1 file changed, 32 insertions(+), 1 deletion(-)

diff --git a/test/mock/index.js b/test/mock/index.js
index f137a822d..3e5e7973e 100644
--- a/test/mock/index.js
+++ b/test/mock/index.js
@@ -1,5 +1,34 @@
 var Server = require('./lib/server');
 
+const cleanup = (servers, spy, callback) => {
+  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
  */
@@ -7,5 +36,7 @@ module.exports = {
   createServer: function(port, host, options) {
     options = options || {};
     return new Server(port, host, options).start();
-  }
+  },
+
+  cleanup: cleanup
 };