1 files changed, 48 insertions, 18 deletions
diff --git a/kamon-spray/src/test/scala/kamon/ServerRequestTracingSpec.scala b/kamon-spray/src/test/scala/kamon/ServerRequestTracingSpec.scala
index 4cff38be..d789042e 100644
--- a/kamon-spray/src/test/scala/kamon/ServerRequestTracingSpec.scala
+++ b/kamon-spray/src/test/scala/kamon/ServerRequestTracingSpec.scala
@@ -13,37 +13,67 @@ import kamon.trace.Trace
 import kamon.Kamon.Extension
 import kamon.trace.UowTracing.{Finish, Start}
 
-class ServerRequestTracingSpec extends TestKit(ActorSystem("server-request-tracing-spec")) with WordSpecLike with RequestBuilding {
+class ServerRequestTracingSpec extends TestKit(ActorSystem("server-request-tracing-spec")) with WordSpecLike with RequestBuilding with TestServer {
 
   "the spray server request tracing instrumentation" should {
-    "start tracing a request when entering the server and close it when responding" in new TestServer {
-      client(Get(s"http://127.0.0.1:$port/"))
+    "trace a request start/finish sequence when proper TraceContext is received" in {
+      send {
+        Get(s"http://127.0.0.1:$port/ok")
+      }
 
       within(5 seconds) {
-        val traceId = expectMsgPF() { case Start(id) => id}
+        val traceId = expectMsgPF() { case Start(id, _) => id}
         expectMsgPF() { case Finish(traceId) => }
       }
     }
-  }
 
+    "finish a request even if no TraceContext is received in the response" in {
+      send {
+        Get(s"http://127.0.0.1:$port/clearcontext")
+      }
 
+      within(5 seconds) {
+        val traceId = expectMsgPF() { case Start(id, _) => id}
+        expectMsgPF() { case Finish(traceId) => }
+      }
+    }
 
-  trait TestServer extends SimpleRoutingApp {
+    "give a initial transaction name using the method and path from the request" in {
+      send {
+        Get(s"http://127.0.0.1:$port/accounts")
+      }
 
-    // Nasty, but very helpful for tests.
-    AkkaExtensionSwap.swap(system, Trace, new Extension {
-      def manager: ActorRef = testActor
-    })
+      within(5 seconds) {
+        expectMsgPF() { case Start(_, "GET: /accounts") => }
+      }
+    }
+  }
+}
 
-    implicit val timeout = Timeout(20 seconds)
-    val port: Int = Await.result(
-        startServer(interface = "127.0.0.1", port = 0)(
-          get {
-            complete("ok")
+trait TestServer extends SimpleRoutingApp {
+  self: TestKit =>
+
+  // Nasty, but very helpful for tests.
+  AkkaExtensionSwap.swap(system, Trace, new Extension {
+    def manager: ActorRef = testActor
+  })
+
+  implicit val timeout = Timeout(20 seconds)
+  val port: Int = Await.result(
+    startServer(interface = "127.0.0.1", port = 0)(
+      get {
+        path("ok") {
+          complete("ok")
+        } ~
+          path("clearcontext"){
+            complete {
+              Trace.clear
+              "ok"
+            }
           }
-        ), timeout.duration).localAddress.getPort
+      }
+    ), timeout.duration).localAddress.getPort
 
-    val client = sendReceive(system, system.dispatcher, timeout)
+  val send = sendReceive(system, system.dispatcher, timeout)
 
-  }
 }