From 6e5b95be319626c2ee1e6c0f666c1e22ae387fe1 Mon Sep 17 00:00:00 2001
From: Erik <nuclear.ace@gmail.com>
Date: Fri, 4 Dec 2015 10:03:01 -0500
Subject: [PATCH] fix nuclearace/Socket.IO-Client-Swift#97

---
 SocketIO-MacTests/SocketSideEffectTest.swift | 30 +++++++++++++++
 Source/SocketIOClient.swift                  | 39 ++++++++++++++------
 2 files changed, 58 insertions(+), 11 deletions(-)

diff --git a/SocketIO-MacTests/SocketSideEffectTest.swift b/SocketIO-MacTests/SocketSideEffectTest.swift
index 4f35516..84d391f 100644
--- a/SocketIO-MacTests/SocketSideEffectTest.swift
+++ b/SocketIO-MacTests/SocketSideEffectTest.swift
@@ -57,6 +57,36 @@ class SocketSideEffectTest: XCTestCase {
         waitForExpectationsWithTimeout(3, handler: nil)
     }
     
+    func testHandleOnceEvent() {
+        let expectation = expectationWithDescription("handled event")
+        socket.once("test") {data, ack in
+            XCTAssertEqual(data[0] as? String, "hello world")
+            XCTAssertEqual(self.socket.testHandlers.count, 0)
+            expectation.fulfill()
+        }
+        
+        socket.parseSocketMessage("2[\"test\",\"hello world\"]")
+        waitForExpectationsWithTimeout(3, handler: nil)
+    }
+    
+    func testOffWithEvent() {
+        socket.on("test") {data, ack in }
+        XCTAssertEqual(socket.testHandlers.count, 1)
+        socket.on("test") {data, ack in }
+        XCTAssertEqual(socket.testHandlers.count, 2)
+        socket.off("test")
+        XCTAssertEqual(socket.testHandlers.count, 0)
+    }
+    
+    func testOffWithId() {
+        let handler = socket.on("test") {data, ack in }
+        XCTAssertEqual(socket.testHandlers.count, 1)
+        socket.on("test") {data, ack in }
+        XCTAssertEqual(socket.testHandlers.count, 2)
+        socket.off(id: handler)
+        XCTAssertEqual(socket.testHandlers.count, 1)
+    }
+    
     func testHandleBinaryEvent() {
         let expectation = expectationWithDescription("handled binary event")
         socket.on("test") {data, ack in
diff --git a/Source/SocketIOClient.swift b/Source/SocketIOClient.swift
index 51dcb3f..3fef0e3 100644
--- a/Source/SocketIOClient.swift
+++ b/Source/SocketIOClient.swift
@@ -389,39 +389,44 @@ public final class SocketIOClient: NSObject, SocketEngineClient {
         handlers = handlers.filter { $0.event != event }
     }
     
+    /**
+    Removes a handler with the specified UUID gotten from an `on` or `once`
+    */
+    public func off(id id: NSUUID) {
+        DefaultSocketLogger.Logger.log("Removing handler with id: %@", type: logType, args: id)
+        
+        handlers = handlers.filter { $0.id != id }
+    }
+    
     /**
      Adds a handler for an event.
      */
-    public func on(event: String, callback: NormalCallback) {
+    public func on(event: String, callback: NormalCallback) -> NSUUID {
         DefaultSocketLogger.Logger.log("Adding handler for event: %@", type: logType, args: event)
         
         let handler = SocketEventHandler(event: event, id: NSUUID(), callback: callback)
         handlers.append(handler)
+        
+        return handler.id
     }
     
     /**
      Adds a single-use handler for an event.
      */
-    public func once(event: String, callback: NormalCallback) {
+    public func once(event: String, callback: NormalCallback) -> NSUUID {
         DefaultSocketLogger.Logger.log("Adding once handler for event: %@", type: logType, args: event)
         
         let id = NSUUID()
         
         let handler = SocketEventHandler(event: event, id: id) {[weak self] data, ack in
             guard let this = self else {return}
-            this.handlers = this.handlers.filter {$0.id != id}
+            this.off(id: id)
             callback(data, ack)
         }
         
         handlers.append(handler)
-    }
-    
-    /**
-     Removes all handlers.
-     Can be used after disconnecting to break any potential remaining retain cycles.
-     */
-    public func removeAllHandlers() {
-        handlers.removeAll(keepCapacity: false)
+        
+        return handler.id
     }
     
     /**
@@ -457,6 +462,14 @@ public final class SocketIOClient: NSObject, SocketEngineClient {
         tryReconnect()
     }
     
+    /**
+     Removes all handlers.
+     Can be used after disconnecting to break any potential remaining retain cycles.
+     */
+    public func removeAllHandlers() {
+        handlers.removeAll(keepCapacity: false)
+    }
+    
     private func tryReconnect() {
         if reconnectTimer == nil {
             DefaultSocketLogger.Logger.log("Starting reconnect", type: logType)
@@ -495,6 +508,10 @@ public final class SocketIOClient: NSObject, SocketEngineClient {
 
 // Test extensions
 extension SocketIOClient {
+    var testHandlers: [SocketEventHandler] {
+        return handlers
+    }
+    
     func setTestable() {
         status = .Connected
     }