From 41d30040be71068fcc7bd82e84df1315b497deaf Mon Sep 17 00:00:00 2001
From: Luigi Pinca <luigipinca@gmail.com>
Date: Wed, 3 Oct 2018 11:54:20 +0200
Subject: [PATCH] [feature] Make `WebSocketServer#close()` emit the `'close'`
 event

Fixes #1375
---
 doc/ws.md                     |  6 ++++++
 lib/websocket-server.js       | 19 +++++++++++++++++--
 test/websocket-server.test.js |  9 ++++++++-
 3 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/doc/ws.md b/doc/ws.md
index 1a19229ff..5c36068b3 100644
--- a/doc/ws.md
+++ b/doc/ws.md
@@ -101,6 +101,12 @@ to the threshold. This determines if compression is used for the entire message.
 `callback` will be added as a listener for the `listening` event on the HTTP
 server when not operating in "noServer" mode.
 
+### Event: 'close'
+
+Emitted when the server closes. This event depends on the `'close'` event of
+HTTP server only when it is created internally. In all other cases, the event
+is emitted independently.
+
 ### Event: 'connection'
 
 - `socket` {WebSocket}
diff --git a/lib/websocket-server.js b/lib/websocket-server.js
index 70513edf2..709f38852 100644
--- a/lib/websocket-server.js
+++ b/lib/websocket-server.js
@@ -112,6 +112,8 @@ class WebSocketServer extends EventEmitter {
    * @public
    */
   close (cb) {
+    if (cb) this.once('close', cb);
+
     //
     // Terminate all associated clients.
     //
@@ -128,10 +130,13 @@ class WebSocketServer extends EventEmitter {
       //
       // Close the http server if it was internally created.
       //
-      if (this.options.port != null) return server.close(cb);
+      if (this.options.port != null) {
+        server.close(() => this.emit('close'));
+        return;
+      }
     }
 
-    if (cb) cb();
+    process.nextTick(emitClose, this);
   }
 
   /**
@@ -317,6 +322,16 @@ function addListeners (server, map) {
   };
 }
 
+/**
+ * Emit a `'close'` event on an `EventEmitter`.
+ *
+ * @param {EventEmitter} server The event emitter
+ * @private
+ */
+function emitClose (server) {
+  server.emit('close');
+}
+
 /**
  * Handle premature socket errors.
  *
diff --git a/test/websocket-server.test.js b/test/websocket-server.test.js
index 45d287978..46af9f10b 100644
--- a/test/websocket-server.test.js
+++ b/test/websocket-server.test.js
@@ -230,6 +230,13 @@ describe('WebSocketServer', function () {
         });
       });
     });
+
+    it("emits the 'close' event", function (done) {
+      const wss = new WebSocket.Server({ noServer: true });
+
+      wss.on('close', done);
+      wss.close();
+    });
   });
 
   describe('#clients', function () {
@@ -689,7 +696,7 @@ describe('WebSocketServer', function () {
       });
     });
 
-    it('emits the `headers` event', function (done) {
+    it("emits the 'headers' event", function (done) {
       const wss = new WebSocket.Server({ port: 0 }, () => {
         const ws = new WebSocket(`ws://localhost:${wss.address().port}`);