From 619701df531134edabbdac2f214c436b669705a3 Mon Sep 17 00:00:00 2001
From: Atrue <wtf.ozu@gmail.com>
Date: Thu, 9 Feb 2023 11:37:46 +0400
Subject: [PATCH] add off method to remove the listener

To have consistency with the `on` method
---
 README.md                    |  1 +
 index.d.ts                   |  3 ++-
 src/helpers/proxy-factory.js |  9 +++++++--
 src/server.js                |  7 +++++++
 tests/unit/server.test.js    | 14 +++++++++++---
 5 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index 7c55495a..b9cb6547 100644
--- a/README.md
+++ b/README.md
@@ -107,6 +107,7 @@ const mockServer = new Server('ws://localhost:8080');
 mockServer.on('connection', socket => {
   socket.on('message', () => {});
   socket.on('close', () => {});
+  socket.on('error', () => {});
 
   socket.send('message');
   socket.close();
diff --git a/index.d.ts b/index.d.ts
index 45ad4374..70b3b6d9 100644
--- a/index.d.ts
+++ b/index.d.ts
@@ -17,7 +17,6 @@ declare module 'mock-socket' {
     close: () => void;
     error: (err: Error) => void;
     message: (message: string | Blob | ArrayBuffer | ArrayBufferView) => void;
-    open: () => void;
   }
 
   //
@@ -56,6 +55,7 @@ declare module 'mock-socket' {
     target: WebSocket;
     close(options?: CloseOptions): void;
     on<K extends keyof WebSocketCallbackMap>(type: K, callback: WebSocketCallbackMap[K]): void;
+    off<K extends keyof WebSocketCallbackMap>(type: K, callback: WebSocketCallbackMap[K]): void;
   }
 
   class Server extends EventTarget {
@@ -68,6 +68,7 @@ declare module 'mock-socket' {
     restoreWebsocket(): void;
 
     on(type: string, callback: (socket: Client) => void): void;
+    off(type: string, callback: (socket: Client) => void): void;
     close(options?: CloseOptions): void;
     emit(event: string, data: any, options?: EmitOptions): void;
 
diff --git a/src/helpers/proxy-factory.js b/src/helpers/proxy-factory.js
index 83766c26..f1b8d094 100644
--- a/src/helpers/proxy-factory.js
+++ b/src/helpers/proxy-factory.js
@@ -36,10 +36,15 @@ export default function proxyFactory(target) {
         };
       }
 
+      const toSocketName = type => (type === 'message' ? `server::${type}` : type);
       if (prop === 'on') {
         return function onWrapper(type, cb) {
-          type = type === 'message' ? `server::${type}` : type;
-          target.addEventListener(type, cb);
+          target.addEventListener(toSocketName(type), cb);
+        };
+      }
+      if (prop === 'off') {
+        return function offWrapper(type, cb) {
+          target.removeEventListener(toSocketName(type), cb);
         };
       }
 
diff --git a/src/server.js b/src/server.js
index e28a037e..a09552a0 100644
--- a/src/server.js
+++ b/src/server.js
@@ -91,6 +91,13 @@ class Server extends EventTarget {
     this.addEventListener(type, callback);
   }
 
+  /*
+  * Remove event listener
+  */
+  off(type, callback) {
+    this.removeEventListener(type, callback);
+  }
+
   /*
    * Closes the connection and triggers the onclose method of all listening
    * websockets. After that it removes itself from the urlMap so another server
diff --git a/tests/unit/server.test.js b/tests/unit/server.test.js
index e33a49d5..bc1e3e87 100644
--- a/tests/unit/server.test.js
+++ b/tests/unit/server.test.js
@@ -20,15 +20,23 @@ test('that after creating a server it is added to the network bridge', t => {
   t.deepEqual(networkBridge.urlMap, {}, 'the urlMap was cleared after the close call');
 });
 
-test('that callback functions can be added to the listeners object', t => {
+test('that callback functions can be added and removed from the listeners object', t => {
   const myServer = new Server('ws://not-real/');
 
-  myServer.on('message', () => {});
-  myServer.on('close', () => {});
+  const onMessage = () => {};
+  const onError = () => {};
+  myServer.on('message', onMessage);
+  myServer.on('close', onError);
 
   t.is(myServer.listeners.message.length, 1);
   t.is(myServer.listeners.close.length, 1);
 
+  myServer.off('message', onMessage);
+  myServer.off('close', onError);
+
+  t.is(myServer.listeners.message.length, 0);
+  t.is(myServer.listeners.close.length, 0);
+
   myServer.close();
 });