halfway to our own NewRelic Agent

2 files changed, 50 insertions, 19 deletions

diff --git a/kamon-spray/src/main/scala/spray/can/server/ServerRequestTracing.scala b/kamon-spray/src/main/scala/spray/can/server/ServerRequestTracing.scala
index d5e21f35..8055cf6b 100644
--- a/kamon-spray/src/main/scala/spray/can/server/ServerRequestTracing.scala
+++ b/kamon-spray/src/main/scala/spray/can/server/ServerRequestTracing.scala
@@ -42,8 +42,9 @@ class ServerRequestTracing {
           publishWarning(s"Different ids when trying to close a Trace, original: [$original] - incoming: [$incoming]")
 
         case Some(_) => // nothing to do here.
-          
+
         case None =>
+          original.finish
           publishWarning(s"Trace context not present while closing the Trace: [$original]")
       }
     }
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)
 
-  }
 }