Sketch out a standard `cask.Logger` interface and standardize dependency injection via case class implicits

29 files changed, 87 insertions, 42 deletions

diff --git a/build.sc b/build.sc
index 69fae11..926f049 100644
--- a/build.sc
+++ b/build.sc
@@ -29,6 +29,8 @@ object cask extends ScalaModule with PublishModule {
     ivy"org.scala-lang:scala-reflect:${scalaVersion()}",
     ivy"io.undertow:undertow-core:2.0.13.Final",
     ivy"com.lihaoyi::upickle:0.7.5",
+    ivy"com.lihaoyi::sourcecode:0.1.7",
+    ivy"com.lihaoyi::pprint:0.5.5"
   )
   def compileIvyDeps = Agg(ivy"com.lihaoyi::acyclic:0.2.0")
   def scalacOptions = Seq("-P:acyclic:force")
diff --git a/cask/src/cask/endpoints/WebSocketEndpoint.scala b/cask/src/cask/endpoints/WebSocketEndpoint.scala
index 6728581..a836321 100644
--- a/cask/src/cask/endpoints/WebSocketEndpoint.scala
+++ b/cask/src/cask/endpoints/WebSocketEndpoint.scala
@@ -4,6 +4,7 @@ import java.nio.ByteBuffer
 
 import cask.internal.{BatchActor, Router}
 import cask.model.Request
+import cask.util.Logger
 import io.undertow.websockets.WebSocketConnectionCallback
 import io.undertow.websockets.core.{AbstractReceiveListener, BufferedBinaryMessage, BufferedTextMessage, CloseMessage, WebSocketChannel, WebSockets}
 import io.undertow.websockets.spi.WebSocketHttpExchange
@@ -31,7 +32,8 @@ class websocket(val path: String, override val subpath: Boolean = false)
   def wrapPathSegment(s: String): Seq[String] = Seq(s)
 }
 
-case class WsHandler(f: WsChannelActor => BatchActor[WsActor.Event])(implicit ec: ExecutionContext)
+case class WsHandler(f: WsChannelActor => BatchActor[WsActor.Event])
+                    (implicit ec: ExecutionContext, log: Logger)
 extends WebsocketResult with WebSocketConnectionCallback {
    def onConnect(exchange: WebSocketHttpExchange, channel: WebSocketChannel): Unit = {
      val actor = f(new WsChannelActor(channel))
@@ -67,7 +69,8 @@ extends WebsocketResult with WebSocketConnectionCallback {
   }
 }
 
-class WsChannelActor(channel: WebSocketChannel)(implicit ec: ExecutionContext)
+class WsChannelActor(channel: WebSocketChannel)
+                    (implicit ec: ExecutionContext, log: Logger)
 extends BatchActor[WsActor.Event]{
   def run(items: Seq[WsActor.Event]): Unit = items.foreach{
     case WsActor.Text(value) => WebSockets.sendTextBlocking(value, channel)
@@ -79,7 +82,7 @@ extends BatchActor[WsActor.Event]{
 }
 
 case class WsActor(handle: PartialFunction[WsActor.Event, Unit])
-                  (implicit ec: ExecutionContext)
+                  (implicit ec: ExecutionContext, log: Logger)
 extends BatchActor[WsActor.Event]{
   def run(items: Seq[WsActor.Event]): Unit = {
     items.foreach(handle.applyOrElse(_, (x: WsActor.Event) => ()))
diff --git a/cask/src/cask/internal/BatchActor.scala b/cask/src/cask/internal/BatchActor.scala
index 1566a18..60b5f57 100644
--- a/cask/src/cask/internal/BatchActor.scala
+++ b/cask/src/cask/internal/BatchActor.scala
@@ -1,5 +1,7 @@
 package cask.internal
 
+import cask.util.Logger
+
 import scala.collection.mutable
 import scala.concurrent.ExecutionContext
 
@@ -8,7 +10,8 @@ import scala.concurrent.ExecutionContext
  * of queued items. `run` handles items in batches, to allow for batch
  * processing optimizations to be used where relevant.
  */
-abstract class BatchActor[T]()(implicit ec: ExecutionContext) {
+abstract class BatchActor[T]()(implicit ec: ExecutionContext,
+                               log: Logger) {
   def run(items: Seq[T]): Unit
 
   private val queue = new mutable.Queue[T]()
@@ -24,7 +27,7 @@ abstract class BatchActor[T]()(implicit ec: ExecutionContext) {
   def runWithItems(): Unit = {
     val items = synchronized(queue.dequeueAll(_ => true))
     try run(items)
-    catch{case e: Throwable => e.printStackTrace()}
+    catch{case e: Throwable => log.exception(e)}
     synchronized{
       if (queue.nonEmpty) ec.execute(() => runWithItems())
       else{
diff --git a/cask/src/cask/main/Main.scala b/cask/src/cask/main/Main.scala
index a5b4dc4..94f7f41 100644
--- a/cask/src/cask/main/Main.scala
+++ b/cask/src/cask/main/Main.scala
@@ -13,7 +13,7 @@ import io.undertow.util.HttpString
   * A combination of [[cask.Main]] and [[cask.Routes]], ideal for small
   * one-file web applications.
   */
-class MainRoutes extends BaseMain with Routes{
+class MainRoutes extends Main with Routes{
   def allRoutes = Seq(this)
 }
 
@@ -24,16 +24,13 @@ class MainRoutes extends BaseMain with Routes{
   * serve, and override various properties on [[Main]] in order to configure
   * application-wide properties.
   */
-class Main(servers0: Routes*) extends BaseMain{
-  def allRoutes = servers0.toSeq
-}
-abstract class BaseMain{
+abstract class Main{
   def mainDecorators = Seq.empty[cask.main.RawDecorator]
   def allRoutes: Seq[Routes]
   def port: Int = 8080
   def host: String = "localhost"
   def debugMode: Boolean = true
-
+  implicit def log = new cask.util.Logger.Console()
   lazy val routeList = for{
     routes <- allRoutes
     route <- routes.caskMetadata.value.map(x => x: EndpointMetadata[_])
@@ -126,11 +123,16 @@ abstract class BaseMain{
                           routes: Routes,
                           metadata: EndpointMetadata[_],
                           e: Router.Result.Error) = {
+    e match {
+      case e: Router.Result.Error.Exception => log.exception(e.t)
+      case _ => // do nothing
+    }
     val statusCode = e match {
       case _: Router.Result.Error.Exception => 500
       case _: Router.Result.Error.InvalidArguments => 400
       case _: Router.Result.Error.MismatchedArguments => 400
     }
+
     val str =
       if (!debugMode) s"Error $statusCode: ${Status.codesToStatus(statusCode).reason}"
       else ErrorMsgs.formatInvokeError(
@@ -138,7 +140,7 @@ abstract class BaseMain{
         metadata.entryPoint.asInstanceOf[EntryPoint[cask.main.Routes, _]],
         e
       )
-    println(str)
+
     Response(str, statusCode = statusCode)
 
   }
diff --git a/cask/src/cask/main/Routes.scala b/cask/src/cask/main/Routes.scala
index be33581..9be9f50 100644
--- a/cask/src/cask/main/Routes.scala
+++ b/cask/src/cask/main/Routes.scala
@@ -1,9 +1,7 @@
 package cask.main
 
-
 import language.experimental.macros
 
-
 trait Routes{
 
   def decorators = Seq.empty[cask.main.RawDecorator]
@@ -15,5 +13,6 @@ trait Routes{
   protected[this] def initialize()(implicit routes: RoutesEndpointsMetadata[this.type]): Unit = {
     metadata0 = routes
   }
-}
 
+  def log: cask.util.Logger
+}
diff --git a/cask/src/cask/package.scala b/cask/src/cask/package.scala
index f72e14b..93c1161 100644
--- a/cask/src/cask/package.scala
+++ b/cask/src/cask/package.scala
@@ -1,3 +1,5 @@
+import cask.util.Logger
+
 package object cask {
   // model
   type Response[T] = model.Response[T]
@@ -45,4 +47,8 @@ package object cask {
   type WsActor = cask.endpoints.WsActor
   val WsActor = cask.endpoints.WsActor
   type WsChannelActor = cask.endpoints.WsChannelActor
+
+  // util
+  type Logger = util.Logger
+  val Logger = util.Logger
 }
diff --git a/cask/src/cask/util/Logger.scala b/cask/src/cask/util/Logger.scala
new file mode 100644
index 0000000..8dc3156
--- /dev/null
+++ b/cask/src/cask/util/Logger.scala
@@ -0,0 +1,18 @@
+package cask.util
+
+import sourcecode.{File, Line, Text}
+
+trait Logger {
+  def exception(t: Throwable): Unit
+
+  def debug(t: sourcecode.Text[Any])(implicit f: sourcecode.File, line: sourcecode.Line): Unit
+}
+object Logger{
+  class Console() extends Logger{
+    def exception(t: Throwable): Unit = t.printStackTrace()
+
+    def debug(t: Text[Any])(implicit f: File, line: Line): Unit = {
+      println(f.value.split('/').last + ":" + line + " " + t.source + " " + pprint.apply(t.value))
+    }
+  }
+}
\ No newline at end of file
diff --git a/example/compress/app/test/src/ExampleTests.scala b/example/compress/app/test/src/ExampleTests.scala
index cd4d8d0..68c249f 100644
--- a/example/compress/app/test/src/ExampleTests.scala
+++ b/example/compress/app/test/src/ExampleTests.scala
@@ -4,7 +4,7 @@ import io.undertow.Undertow
 import utest._
 
 object ExampleTests extends TestSuite{
-  def withServer[T](example: cask.main.BaseMain)(f: String => T): T = {
+  def withServer[T](example: cask.main.Main)(f: String => T): T = {
     val server = Undertow.builder
       .addHttpListener(8080, "localhost")
       .setHandler(example.defaultHandler)
diff --git a/example/compress2/app/src/Compress2.scala b/example/compress2/app/src/Compress2.scala
index 1a4cf69..b0deaa0 100644
--- a/example/compress2/app/src/Compress2.scala
+++ b/example/compress2/app/src/Compress2.scala
@@ -1,6 +1,8 @@
 package app
 
-object Compress2 extends cask.Routes{
+import cask.util.Logger
+
+case class Compress2()(implicit val log: Logger) extends cask.Routes{
   override def decorators = Seq(new cask.decorators.compress())
 
   @cask.get("/")
@@ -11,4 +13,6 @@ object Compress2 extends cask.Routes{
   initialize()
 }
 
-object Compress2Main extends cask.Main(Compress2)
\ No newline at end of file
+object Compress2Main extends cask.Main{
+  val allRoutes = Seq(Compress2())
+}
diff --git a/example/compress2/app/test/src/ExampleTests.scala b/example/compress2/app/test/src/ExampleTests.scala
index 9f9b4b3..fdfedda 100644
--- a/example/compress2/app/test/src/ExampleTests.scala
+++ b/example/compress2/app/test/src/ExampleTests.scala
@@ -4,7 +4,7 @@ import io.undertow.Undertow
 import utest._
 
 object ExampleTests extends TestSuite{
-  def withServer[T](example: cask.main.BaseMain)(f: String => T): T = {
+  def withServer[T](example: cask.main.Main)(f: String => T): T = {
     val server = Undertow.builder
       .addHttpListener(8080, "localhost")
       .setHandler(example.defaultHandler)
diff --git a/example/compress3/app/src/Compress3.scala b/example/compress3/app/src/Compress3.scala
index 4d4df99..95bd851 100644
--- a/example/compress3/app/src/Compress3.scala
+++ b/example/compress3/app/src/Compress3.scala
@@ -1,6 +1,8 @@
 package app
 
-object Compress3 extends cask.Routes{
+import cask.util.Logger
+
+case class Compress3()(implicit val log: Logger) extends cask.Routes{
 
   @cask.get("/")
   def hello() = {
@@ -10,6 +12,7 @@ object Compress3 extends cask.Routes{
   initialize()
 }
 
-object Compress3Main extends cask.Main(Compress3){
+object Compress3Main extends cask.Main{
   override def mainDecorators = Seq(new cask.decorators.compress())
+  val allRoutes = Seq(Compress3())
 }
\ No newline at end of file
diff --git a/example/compress3/app/test/src/ExampleTests.scala b/example/compress3/app/test/src/ExampleTests.scala
index 88e1cbf..2bd8b31 100644
--- a/example/compress3/app/test/src/ExampleTests.scala
+++ b/example/compress3/app/test/src/ExampleTests.scala
@@ -4,7 +4,7 @@ import io.undertow.Undertow
 import utest._
 
 object ExampleTests extends TestSuite{
-  def withServer[T](example: cask.main.BaseMain)(f: String => T): T = {
+  def withServer[T](example: cask.main.Main)(f: String => T): T = {
     val server = Undertow.builder
       .addHttpListener(8080, "localhost")
       .setHandler(example.defaultHandler)
@@ -20,8 +20,9 @@ object ExampleTests extends TestSuite{
     test("Compress3Main") - withServer(Compress3Main){ host =>
       val expected = "Hello World! Hello World! Hello World!"
       requests.get(s"$host").text() ==> expected
+      val compressed = requests.get(s"$host", autoDecompress = false).text()
       assert(
-        requests.get(s"$host", autoDecompress = false).text().length < expected.length
+        compressed.length < expected.length
       )
 
     }
diff --git a/example/cookies/app/test/src/ExampleTests.scala b/example/cookies/app/test/src/ExampleTests.scala
index 150f59a..5954e35 100644
--- a/example/cookies/app/test/src/ExampleTests.scala
+++ b/example/cookies/app/test/src/ExampleTests.scala
@@ -4,7 +4,7 @@ import io.undertow.Undertow
 import utest._
 
 object ExampleTests extends TestSuite{
-  def withServer[T](example: cask.main.BaseMain)(f: String => T): T = {
+  def withServer[T](example: cask.main.Main)(f: String => T): T = {
     val server = Undertow.builder
       .addHttpListener(8080, "localhost")
       .setHandler(example.defaultHandler)
diff --git a/example/decorated/app/test/src/ExampleTests.scala b/example/decorated/app/test/src/ExampleTests.scala
index c2c19a6..f6c8d01 100644
--- a/example/decorated/app/test/src/ExampleTests.scala
+++ b/example/decorated/app/test/src/ExampleTests.scala
@@ -4,7 +4,7 @@ import io.undertow.Undertow
 import utest._
 
 object ExampleTests extends TestSuite{
-  def withServer[T](example: cask.main.BaseMain)(f: String => T): T = {
+  def withServer[T](example: cask.main.Main)(f: String => T): T = {
     val server = Undertow.builder
       .addHttpListener(8080, "localhost")
       .setHandler(example.defaultHandler)
diff --git a/example/decorated2/app/test/src/ExampleTests.scala b/example/decorated2/app/test/src/ExampleTests.scala
index 4343097..624b86d 100644
--- a/example/decorated2/app/test/src/ExampleTests.scala
+++ b/example/decorated2/app/test/src/ExampleTests.scala
@@ -4,7 +4,7 @@ import io.undertow.Undertow
 import utest._
 
 object ExampleTests extends TestSuite{
-  def withServer[T](example: cask.main.BaseMain)(f: String => T): T = {
+  def withServer[T](example: cask.main.Main)(f: String => T): T = {
     val server = Undertow.builder
       .addHttpListener(8080, "localhost")
       .setHandler(example.defaultHandler)
diff --git a/example/endpoints/app/test/src/ExampleTests.scala b/example/endpoints/app/test/src/ExampleTests.scala
index 3314f24..a604f2f 100644
--- a/example/endpoints/app/test/src/ExampleTests.scala
+++ b/example/endpoints/app/test/src/ExampleTests.scala
@@ -4,7 +4,7 @@ import io.undertow.Undertow
 import utest._
 
 object ExampleTests extends TestSuite{
-  def withServer[T](example: cask.main.BaseMain)(f: String => T): T = {
+  def withServer[T](example: cask.main.Main)(f: String => T): T = {
     val server = Undertow.builder
       .addHttpListener(8080, "localhost")
       .setHandler(example.defaultHandler)
diff --git a/example/formJsonPost/app/test/src/ExampleTests.scala b/example/formJsonPost/app/test/src/ExampleTests.scala
index 4178497..768220c 100644
--- a/example/formJsonPost/app/test/src/ExampleTests.scala
+++ b/example/formJsonPost/app/test/src/ExampleTests.scala
@@ -4,7 +4,7 @@ import io.undertow.Undertow
 import utest._
 
 object ExampleTests extends TestSuite{
-  def withServer[T](example: cask.main.BaseMain)(f: String => T): T = {
+  def withServer[T](example: cask.main.Main)(f: String => T): T = {
     val server = Undertow.builder
       .addHttpListener(8080, "localhost")
       .setHandler(example.defaultHandler)
diff --git a/example/httpMethods/app/test/src/ExampleTests.scala b/example/httpMethods/app/test/src/ExampleTests.scala
index 30fa87f..40756a5 100644
--- a/example/httpMethods/app/test/src/ExampleTests.scala
+++ b/example/httpMethods/app/test/src/ExampleTests.scala
@@ -4,7 +4,7 @@ import io.undertow.Undertow
 import utest._
 
 object ExampleTests extends TestSuite{
-  def withServer[T](example: cask.main.BaseMain)(f: String => T): T = {
+  def withServer[T](example: cask.main.Main)(f: String => T): T = {
     val server = Undertow.builder
       .addHttpListener(8080, "localhost")
       .setHandler(example.defaultHandler)
diff --git a/example/minimalApplication/app/test/src/ExampleTests.scala b/example/minimalApplication/app/test/src/ExampleTests.scala
index 1cda7a6..4ae6095 100644
--- a/example/minimalApplication/app/test/src/ExampleTests.scala
+++ b/example/minimalApplication/app/test/src/ExampleTests.scala
@@ -4,7 +4,7 @@ import io.undertow.Undertow
 import utest._
 
 object ExampleTests extends TestSuite{
-  def withServer[T](example: cask.main.BaseMain)(f: String => T): T = {
+  def withServer[T](example: cask.main.Main)(f: String => T): T = {
     val server = Undertow.builder
       .addHttpListener(8080, "localhost")
       .setHandler(example.defaultHandler)
diff --git a/example/minimalApplication2/app/src/MinimalApplication2.scala b/example/minimalApplication2/app/src/MinimalApplication2.scala
index e537ea6..41ca55f 100644
--- a/example/minimalApplication2/app/src/MinimalApplication2.scala
+++ b/example/minimalApplication2/app/src/MinimalApplication2.scala
@@ -1,6 +1,8 @@
 package app
 
-object MinimalRoutes extends cask.Routes{
+import cask.util.Logger
+
+case class MinimalRoutes()(implicit val log: Logger) extends cask.Routes{
   @cask.get("/")
   def hello() = {
     "Hello World!"
@@ -13,4 +15,6 @@ object MinimalRoutes extends cask.Routes{
 
   initialize()
 }
-object MinimalMain extends cask.Main(MinimalRoutes)
\ No newline at end of file
+object MinimalRoutesMain extends cask.Main{
+  val allRoutes = Seq(MinimalRoutes())
+}
\ No newline at end of file
diff --git a/example/minimalApplication2/app/test/src/ExampleTests.scala b/example/minimalApplication2/app/test/src/ExampleTests.scala
index 4e5621c..a77c490 100644
--- a/example/minimalApplication2/app/test/src/ExampleTests.scala
+++ b/example/minimalApplication2/app/test/src/ExampleTests.scala
@@ -4,7 +4,7 @@ import io.undertow.Undertow
 import utest._
 
 object ExampleTests extends TestSuite{
-  def withServer[T](example: cask.main.BaseMain)(f: String => T): T = {
+  def withServer[T](example: cask.main.Main)(f: String => T): T = {
     val server = Undertow.builder
       .addHttpListener(8080, "localhost")
       .setHandler(example.defaultHandler)
@@ -17,7 +17,7 @@ object ExampleTests extends TestSuite{
   }
 
   val tests = Tests{
-    test("MinimalApplication2") - withServer(MinimalMain){ host =>
+    test("MinimalApplication2") - withServer(MinimalRoutesMain){ host =>
       val success = requests.get(host)
 
       success.text() ==> "Hello World!"
diff --git a/example/redirectAbort/app/test/src/ExampleTests.scala b/example/redirectAbort/app/test/src/ExampleTests.scala
index a4d149f..d87d24a 100644
--- a/example/redirectAbort/app/test/src/ExampleTests.scala
+++ b/example/redirectAbort/app/test/src/ExampleTests.scala
@@ -4,7 +4,7 @@ import io.undertow.Undertow
 import utest._
 
 object ExampleTests extends TestSuite{
-  def withServer[T](example: cask.main.BaseMain)(f: String => T): T = {
+  def withServer[T](example: cask.main.Main)(f: String => T): T = {
     val server = Undertow.builder
       .addHttpListener(8080, "localhost")
       .setHandler(example.defaultHandler)
diff --git a/example/scalatags/app/test/src/ExampleTests.scala b/example/scalatags/app/test/src/ExampleTests.scala
index 53bc1ea..d4dfbb3 100644
--- a/example/scalatags/app/test/src/ExampleTests.scala
+++ b/example/scalatags/app/test/src/ExampleTests.scala
@@ -4,7 +4,7 @@ import io.undertow.Undertow
 import utest._
 
 object ExampleTests extends TestSuite{
-  def withServer[T](example: cask.main.BaseMain)(f: String => T): T = {
+  def withServer[T](example: cask.main.Main)(f: String => T): T = {
     val server = Undertow.builder
       .addHttpListener(8080, "localhost")
       .setHandler(example.defaultHandler)
diff --git a/example/staticFiles/app/test/src/ExampleTests.scala b/example/staticFiles/app/test/src/ExampleTests.scala
index ebda8a0..8b74d19 100644
--- a/example/staticFiles/app/test/src/ExampleTests.scala
+++ b/example/staticFiles/app/test/src/ExampleTests.scala
@@ -4,7 +4,7 @@ import io.undertow.Undertow
 import utest._
 
 object ExampleTests extends TestSuite{
-  def withServer[T](example: cask.main.BaseMain)(f: String => T): T = {
+  def withServer[T](example: cask.main.Main)(f: String => T): T = {
     val server = Undertow.builder
       .addHttpListener(8080, "localhost")
       .setHandler(example.defaultHandler)
diff --git a/example/todoApi/app/test/src/ExampleTests.scala b/example/todoApi/app/test/src/ExampleTests.scala
index 4e85a8e..f65f083 100644
--- a/example/todoApi/app/test/src/ExampleTests.scala
+++ b/example/todoApi/app/test/src/ExampleTests.scala
@@ -4,7 +4,7 @@ import io.undertow.Undertow
 import utest._
 
 object ExampleTests extends TestSuite{
-  def withServer[T](example: cask.main.BaseMain)(f: String => T): T = {
+  def withServer[T](example: cask.main.Main)(f: String => T): T = {
     val server = Undertow.builder
       .addHttpListener(8080, "localhost")
       .setHandler(example.defaultHandler)
diff --git a/example/twirl/app/test/src/ExampleTests.scala b/example/twirl/app/test/src/ExampleTests.scala
index 2c445fb..0dcc120 100644
--- a/example/twirl/app/test/src/ExampleTests.scala
+++ b/example/twirl/app/test/src/ExampleTests.scala
@@ -4,7 +4,7 @@ import io.undertow.Undertow
 import utest._
 
 object ExampleTests extends TestSuite{
-  def withServer[T](example: cask.main.BaseMain)(f: String => T): T = {
+  def withServer[T](example: cask.main.Main)(f: String => T): T = {
     val server = Undertow.builder
       .addHttpListener(8080, "localhost")
       .setHandler(example.defaultHandler)
diff --git a/example/variableRoutes/app/test/src/ExampleTests.scala b/example/variableRoutes/app/test/src/ExampleTests.scala
index 1755dea..93673d9 100644
--- a/example/variableRoutes/app/test/src/ExampleTests.scala
+++ b/example/variableRoutes/app/test/src/ExampleTests.scala
@@ -4,7 +4,7 @@ import io.undertow.Undertow
 import utest._
 
 object ExampleTests extends TestSuite{
-  def withServer[T](example: cask.main.BaseMain)(f: String => T): T = {
+  def withServer[T](example: cask.main.Main)(f: String => T): T = {
     val server = Undertow.builder
       .addHttpListener(8080, "localhost")
       .setHandler(example.defaultHandler)
diff --git a/example/websockets/app/test/src/ExampleTests.scala b/example/websockets/app/test/src/ExampleTests.scala
index a463824..d4d96da 100644
--- a/example/websockets/app/test/src/ExampleTests.scala
+++ b/example/websockets/app/test/src/ExampleTests.scala
@@ -8,7 +8,7 @@ import utest._
 object ExampleTests extends TestSuite{
 
 
-  def withServer[T](example: cask.main.BaseMain)(f: String => T): T = {
+  def withServer[T](example: cask.main.Main)(f: String => T): T = {
     val server = io.undertow.Undertow.builder
       .addHttpListener(8080, "localhost")
       .setHandler(example.defaultHandler)
diff --git a/example/websockets2/app/test/src/ExampleTests.scala b/example/websockets2/app/test/src/ExampleTests.scala
index 27bff5e..96dfd7e 100644
--- a/example/websockets2/app/test/src/ExampleTests.scala
+++ b/example/websockets2/app/test/src/ExampleTests.scala
@@ -8,7 +8,7 @@ import utest._
 object ExampleTests extends TestSuite{
 
 
-  def withServer[T](example: cask.main.BaseMain)(f: String => T): T = {
+  def withServer[T](example: cask.main.Main)(f: String => T): T = {
     val server = io.undertow.Undertow.builder
       .addHttpListener(8080, "localhost")
       .setHandler(example.defaultHandler)