From 3952eacfaa03af3519cb0d5666d2952bfa5f3de3 Mon Sep 17 00:00:00 2001
From: Mikael Mengistu <mikaelm12@users.noreply.github.com>
Date: Tue, 27 Mar 2018 14:06:56 -0700
Subject: [PATCH] Add ability to remove all handlers with just the method name
 (#1729)

---
 clients/ts/signalr/spec/HubConnection.spec.ts | 88 +++++++++++++++++++
 clients/ts/signalr/src/HubConnection.ts       | 30 ++++---
 2 files changed, 108 insertions(+), 10 deletions(-)

diff --git a/clients/ts/signalr/spec/HubConnection.spec.ts b/clients/ts/signalr/spec/HubConnection.spec.ts
index 1695c09db6..9224c0fa77 100644
--- a/clients/ts/signalr/spec/HubConnection.spec.ts
+++ b/clients/ts/signalr/spec/HubConnection.spec.ts
@@ -271,6 +271,94 @@ describe("HubConnection", () => {
             expect(warnings).toEqual(["No client method with the name 'message' found."]);
         });
 
+        it("all handlers can be unregistered with just the method name", async () => {
+            const connection = new TestConnection();
+            const hubConnection = new HubConnection(connection);
+
+            connection.receiveHandshakeResponse();
+
+            let count = 0;
+            const handler = () => { count++; };
+            const secondHandler = () => { count++; };
+            hubConnection.on("inc", handler);
+            hubConnection.on("inc", secondHandler);
+
+            connection.receive({
+                arguments: [],
+                invocationId: "0",
+                nonblocking: true,
+                target: "inc",
+                type: MessageType.Invocation,
+            });
+
+            hubConnection.off("inc");
+
+            connection.receive({
+                arguments: [],
+                invocationId: "0",
+                nonblocking: true,
+                target: "inc",
+                type: MessageType.Invocation,
+            });
+
+            expect(count).toBe(2);
+        });
+
+        it("a single handler can be unregistered with the method name and handler", async () => {
+            const connection = new TestConnection();
+            const hubConnection = new HubConnection(connection);
+
+            connection.receiveHandshakeResponse();
+
+            let count = 0;
+            const handler = () => { count++; };
+            const secondHandler = () => { count++; };
+            hubConnection.on("inc", handler);
+            hubConnection.on("inc", secondHandler);
+
+            connection.receive({
+                arguments: [],
+                invocationId: "0",
+                nonblocking: true,
+                target: "inc",
+                type: MessageType.Invocation,
+            });
+
+            hubConnection.off("inc", handler);
+
+            connection.receive({
+                arguments: [],
+                invocationId: "0",
+                nonblocking: true,
+                target: "inc",
+                type: MessageType.Invocation,
+            });
+
+            expect(count).toBe(3);
+        });
+
+        it("can't register the same handler multiple times", async () => {
+            const connection = new TestConnection();
+            const hubConnection = new HubConnection(connection);
+
+            connection.receiveHandshakeResponse();
+
+            let count = 0;
+            const handler = () => { count++; };
+            hubConnection.on("inc", handler);
+            hubConnection.on("inc", handler);
+
+            connection.receive({
+                arguments: [],
+                invocationId: "0",
+                nonblocking: true,
+                target: "inc",
+                type: MessageType.Invocation,
+            });
+
+            expect(count).toBe(1);
+        });
+
         it("callback invoked when servers invokes a method on the client", async () => {
             const connection = new TestConnection();
             const hubConnection = new HubConnection(connection, commonOptions);
diff --git a/clients/ts/signalr/src/HubConnection.ts b/clients/ts/signalr/src/HubConnection.ts
index 72c01d31db..8fd97c3d13 100644
--- a/clients/ts/signalr/src/HubConnection.ts
+++ b/clients/ts/signalr/src/HubConnection.ts
@@ -306,8 +306,8 @@ export class HubConnection {
         return p;
     }
 
-    public on(methodName: string, method: (...args: any[]) => void) {
-        if (!methodName || !method) {
+    public on(methodName: string, newMethod: (...args: any[]) => void) {
+        if (!methodName || !newMethod) {
             return;
         }
 
@@ -316,11 +316,16 @@ export class HubConnection {
             this.methods[methodName] = [];
         }
 
-        this.methods[methodName].push(method);
+        // Preventing adding the same handler multiple times.
+        if (this.methods[methodName].indexOf(newMethod) !== -1) {
+            return;
+         }
+
+        this.methods[methodName].push(newMethod);
     }
 
-    public off(methodName: string, method: (...args: any[]) => void) {
-        if (!methodName || !method) {
+    public off(methodName: string, method?: (...args: any[]) => void) {
+        if (!methodName) {
             return;
         }
 
@@ -329,13 +334,18 @@ export class HubConnection {
         if (!handlers) {
             return;
         }
-        const removeIdx = handlers.indexOf(method);
-        if (removeIdx !== -1) {
-            handlers.splice(removeIdx, 1);
-            if (handlers.length === 0) {
-                delete this.methods[methodName];
+        if (method) {
+            const removeIdx = handlers.indexOf(method);
+            if (removeIdx !== -1) {
+                handlers.splice(removeIdx, 1);
+                if (handlers.length === 0) {
+                    delete this.methods[methodName];
+                }
             }
+        } else {
+            delete this.methods[methodName];
         }
+
     }
 
     public onclose(callback: ConnectionClosed) {