From 6086ae0b29db10bd98b423bcd89b43fd2d3e15ba Mon Sep 17 00:00:00 2001
From: Isaac Murchie <isaac@saucelabs.com>
Date: Tue, 21 Jun 2016 12:42:10 -0700
Subject: [PATCH] Make sure the server is closed when we call close

---
 lib/express/server.js               | 12 ++++++++++++
 test/basedriver/driver-e2e-tests.js |  4 ++--
 test/express/server-specs.js        | 21 ++++++++++++++++++++-
 3 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/lib/express/server.js b/lib/express/server.js
index 55c926716..34d5ed59c 100644
--- a/lib/express/server.js
+++ b/lib/express/server.js
@@ -17,6 +17,18 @@ async function server (configureRoutes, port, hostname = null) {
   let app = express();
   let httpServer = http.createServer(app);
 
+  // http.Server.close() only stops new connections, but we need to wait until
+  // all connections are closed and the `close` event is emitted
+  let close = httpServer.close.bind(httpServer);
+  httpServer.close = async () => {
+    return await new Promise((resolve, reject) => {
+      httpServer.on('close', resolve);
+      close((err) => {
+        if (err) reject(err);
+      });
+    });
+  };
+
   return await new Promise((resolve, reject) => {
     httpServer.on('error', (err) => {
       if (err.code === 'EADDRNOTAVAIL') {
diff --git a/test/basedriver/driver-e2e-tests.js b/test/basedriver/driver-e2e-tests.js
index 5374c27e9..9b31e99e6 100644
--- a/test/basedriver/driver-e2e-tests.js
+++ b/test/basedriver/driver-e2e-tests.js
@@ -14,8 +14,8 @@ function baseDriverE2ETests (DriverClass, defaultCaps = {}) {
     before(async () => {
       baseServer = await server(routeConfiguringFunction(d), 8181);
     });
-    after(() => {
-      baseServer.close();
+    after(async () => {
+      await baseServer.close();
     });
 
     describe('session handling', () => {
diff --git a/test/express/server-specs.js b/test/express/server-specs.js
index a3706b523..56f802aba 100644
--- a/test/express/server-specs.js
+++ b/test/express/server-specs.js
@@ -6,6 +6,7 @@ import request from 'request-promise';
 import chai from 'chai';
 import chaiAsPromised from 'chai-as-promised';
 import sinon from 'sinon';
+import B from 'bluebird';
 
 
 chai.should();
@@ -43,11 +44,16 @@ describe('server', () => {
       app.get('/error', () => {
         throw new Error('hahaha');
       });
+      app.get('/pause', async (req, res) => {
+        res.header['content-type'] = 'text/html';
+        await B.delay(1000);
+        res.status(200).send('We have waited!');
+      });
     }
     hwServer = await server(configureRoutes, 8181);
   });
   after(async () => {
-    hwServer.close();
+    await hwServer.close();
     errorStub.restore();
   });
 
@@ -76,4 +82,17 @@ describe('server', () => {
     await server(() => {}, 8181, 'lolcathost').should.be.rejectedWith(/ENOTFOUND/);
     await server(() => {}, 8181, '1.1.1.1').should.be.rejectedWith(/EADDRNOTAVAIL/);
   });
+  it('should wait for the server close connections before finishing closing', async () => {
+    let bodyPromise = request('http://localhost:8181/pause');
+
+    // relinquish control so that we don't close before the request is received
+    await B.delay(100);
+
+    let before = Date.now();
+    await hwServer.close();
+    // expect slightly less than the request waited, since we paused above
+    (Date.now() - before).should.be.above(900);
+
+    (await bodyPromise).should.equal('We have waited!');
+  });
 });