1 files changed, 12 insertions, 10 deletions
diff --git a/cask/src/cask/endpoints/WebSocketEndpoint.scala b/cask/src/cask/endpoints/WebSocketEndpoint.scala
index f747341..842d508 100644
--- a/cask/src/cask/endpoints/WebSocketEndpoint.scala
+++ b/cask/src/cask/endpoints/WebSocketEndpoint.scala
@@ -3,22 +3,24 @@ package cask.endpoints
 import cask.internal.Router
 import cask.model.Request
 import io.undertow.websockets.WebSocketConnectionCallback
-
+import collection.JavaConverters._
 sealed trait WebsocketResult
 object WebsocketResult{
-  implicit class Response(val value: cask.model.Response) extends WebsocketResult
+  implicit class Response[T](value0: cask.model.Response[T])
+                            (implicit f: T => cask.model.Response.Data) extends WebsocketResult{
+    def value = value0.map(f)
+  }
   implicit class Listener(val value: WebSocketConnectionCallback) extends WebsocketResult
 }
 
-class websocket(val path: String, override val subpath: Boolean = false) extends cask.main.BaseEndpoint{
-  type Output = WebsocketResult
+class websocket(val path: String, override val subpath: Boolean = false)
+  extends cask.main.Endpoint[WebsocketResult, Seq[String]]{
   val methods = Seq("websocket")
-  type Input = Seq[String]
   type InputParser[T] = QueryParamReader[T]
-  type Returned = Router.Result[WebsocketResult]
-  def wrapFunction(ctx: Request, delegate: Delegate): Returned = delegate(Map())
-
-  def wrapPathSegment(s: String): Input = Seq(s)
-
+  type OuterReturned = Router.Result[WebsocketResult]
+  def wrapFunction(ctx: Request, delegate: Delegate): OuterReturned = {
+    delegate(WebEndpoint.buildMapFromQueryParams(ctx))
+  }
 
+  def wrapPathSegment(s: String): Seq[String] = Seq(s)
 }