Merge pull request #15 from lihaoyi/lihaoyi-patch-1

Tighten behavior of WebsocketClientImpl

3 files changed, 18 insertions, 9 deletions

diff --git a/cask/util/src-js/cask/util/WebsocketClientImpl.scala b/cask/util/src-js/cask/util/WebsocketClientImpl.scala
index f129893..29ebfa7 100644
--- a/cask/util/src-js/cask/util/WebsocketClientImpl.scala
+++ b/cask/util/src-js/cask/util/WebsocketClientImpl.scala
@@ -7,7 +7,7 @@ abstract class WebsocketClientImpl(url: String) extends WebsocketBase{
   var closed = false
   def connect(): Unit = {
     websocket = new dom.WebSocket(url)
-
+    assert(closed == false)
     websocket.onopen = (e: dom.Event) => onOpen()
     websocket.onmessage = (e: dom.MessageEvent) => onMessage(e.data.asInstanceOf[String])
     websocket.onclose = (e: dom.CloseEvent) => {
@@ -28,6 +28,8 @@ abstract class WebsocketClientImpl(url: String) extends WebsocketBase{
   def onError(ex: Exception): Unit
   def onMessage(value: String): Unit
   def onClose(code: Int, reason: String): Unit
-  def close(): Unit = websocket.close()
+  def close(): Unit = {
+    if (!closed) websocket.close()
+  }
   def isClosed() = closed
-}
\ No newline at end of file
+}
diff --git a/cask/util/src-jvm/cask/util/WebsocketClientImpl.scala b/cask/util/src-jvm/cask/util/WebsocketClientImpl.scala
index 5570356..3344ce4 100644
--- a/cask/util/src-jvm/cask/util/WebsocketClientImpl.scala
+++ b/cask/util/src-jvm/cask/util/WebsocketClientImpl.scala
@@ -4,8 +4,9 @@ import org.java_websocket.handshake.ServerHandshake
 
 abstract class WebsocketClientImpl(url: String) extends WebsocketBase{
   var websocket: Client = null
-
+  var closed = false
   def connect(): Unit = {
+    assert(closed == false)
     websocket = new Client()
     websocket.connect()
   }
@@ -25,15 +26,20 @@ abstract class WebsocketClientImpl(url: String) extends WebsocketBase{
   }
   def onClose(code: Int, reason: String): Unit
   def onError(ex: Exception): Unit
-  def close(): Unit = websocket.close()
+  def close(): Unit = {
+    if (!closed) websocket.close()
+  }
   def isClosed() = websocket.isClosed()
   class Client() extends WebSocketClient(new java.net.URI(url)){
     def onOpen(handshakedata: ServerHandshake) = {
       WebsocketClientImpl.this.onOpen()
     }
     def onMessage(message: String) = WebsocketClientImpl.this.onMessage(message)
-    def onClose(code: Int, reason: String, remote: Boolean) = WebsocketClientImpl.this.onClose(code, reason)
+    def onClose(code: Int, reason: String, remote: Boolean) = {
+      closed = true
+      WebsocketClientImpl.this.onClose(code, reason)
+    }
     def onError(ex: Exception) = WebsocketClientImpl.this.onError(ex)
 
   }
-}
\ No newline at end of file
+}
diff --git a/cask/util/src/cask/util/WsClient.scala b/cask/util/src/cask/util/WsClient.scala
index b0f56d0..f15d7df 100644
--- a/cask/util/src/cask/util/WsClient.scala
+++ b/cask/util/src/cask/util/WsClient.scala
@@ -14,7 +14,8 @@ class WsClient(impl: WebsocketBase)
     case Ws.Close(_, _) => impl.close()
     case Ws.ChannelClosed() => impl.close()
   }
-  def close() = impl.close()
+    
+  def close(code: Int = 1005, reason: String = "") = this.send(Ws.Close(code, reason))
 }
 
 object WsClient{
@@ -52,4 +53,4 @@ object WsClient{
     impl.connect()
     p.future
   }
-}
\ No newline at end of file
+}