different aproach to testing spray instrumentation, wip

4 files changed, 104 insertions, 39 deletions

diff --git a/kamon-core/src/test/resources/META-INF/aop.xml b/kamon-core/src/test/resources/META-INF/aop.xml
new file mode 100644
index 00000000..ea8dce20
--- /dev/null
+++ b/kamon-core/src/test/resources/META-INF/aop.xml
@@ -0,0 +1,13 @@
+<!DOCTYPE aspectj PUBLIC "-//AspectJ//DTD//EN" "http://www.eclipse.org/aspectj/dtd/aspectj.dtd">
+
+<aspectj>
+    <aspects>
+        <!-- Actors -->
+        <aspect name="akka.testkit.TestProbeTracing"/>
+    </aspects>
+
+    <weaver options="-verbose -showWeaveInfo">
+        <include within="akka..*"/>
+    </weaver>
+
+</aspectj>
\ No newline at end of file
diff --git a/kamon-core/src/test/scala/akka/testkit/TestProbeTracing.scala b/kamon-core/src/test/scala/akka/testkit/TestProbeTracing.scala
new file mode 100644
index 00000000..dbbf1246
--- /dev/null
+++ b/kamon-core/src/test/scala/akka/testkit/TestProbeTracing.scala
@@ -0,0 +1,54 @@
+/*
+ * =========================================================================================
+ * Copyright © 2013 the kamon project <http://kamon.io/>
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
+ * except in compliance with the License. You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the
+ * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
+ * either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ * =========================================================================================
+ */
+
+package akka.testkit
+
+import org.aspectj.lang.annotation._
+import kamon.trace.{ ContextAware, Trace }
+import org.aspectj.lang.ProceedingJoinPoint
+import akka.testkit.TestActor.RealMessage
+
+@Aspect
+class TestProbeTracing {
+
+  @DeclareMixin("akka.testkit.TestActor.RealMessage")
+  def mixin: ContextAware = ContextAware.default
+
+  @Pointcut("execution(akka.testkit.TestActor.RealMessage.new(..)) && this(ctx)")
+  def realMessageCreation(ctx: ContextAware): Unit = {}
+
+  @After("realMessageCreation(ctx)")
+  def afterRealMessageCreation(ctx: ContextAware): Unit = {
+    // Necessary to force the initialization of ContextAware at the moment of creation.
+    ctx.traceContext
+  }
+
+  @Pointcut("execution(* akka.testkit.TestProbe.reply(..)) && this(testProbe)")
+  def testProbeReply(testProbe: TestProbe): Unit = {}
+
+  @Around("testProbeReply(testProbe)")
+  def aroundTestProbeReply(pjp: ProceedingJoinPoint, testProbe: TestProbe): Any = {
+    val traceContext = testProbe.lastMessage match {
+      case msg: RealMessage ⇒ msg.asInstanceOf[ContextAware].traceContext
+      case _                ⇒ None
+    }
+
+    Trace.withContext(traceContext) {
+      pjp.proceed
+    }
+  }
+
+}
diff --git a/kamon-spray/src/test/scala/kamon/spray/ClientRequestTracingSpec.scala b/kamon-spray/src/test/scala/kamon/spray/ClientRequestTracingSpec.scala
index eae0c992..e204cd02 100644
--- a/kamon-spray/src/test/scala/kamon/spray/ClientRequestTracingSpec.scala
+++ b/kamon-spray/src/test/scala/kamon/spray/ClientRequestTracingSpec.scala
@@ -14,13 +14,13 @@ class ClientRequestTracingSpec extends TestKit(ActorSystem("server-request-traci
   "the client instrumentation" should {
     "record segments for a client http request" in {
 
-      Trace.start("record-segments", None)(system)
+      /*      Trace.start("record-segments", None)(system)
 
       send {
         Get(s"http://127.0.0.1:$port/ok")
 
         // We don't care about the response, just make sure we finish the Trace after the response has been received.
-      } map (rsp ⇒ Trace.finish())
+      } map (rsp ⇒ Trace.finish())*/
 
       val trace = expectMsgType[UowTrace]
       println(trace.segments)
diff --git a/kamon-spray/src/test/scala/kamon/spray/ServerRequestTracingSpec.scala b/kamon-spray/src/test/scala/kamon/spray/ServerRequestTracingSpec.scala
index 177e23c4..904047ed 100644
--- a/kamon-spray/src/test/scala/kamon/spray/ServerRequestTracingSpec.scala
+++ b/kamon-spray/src/test/scala/kamon/spray/ServerRequestTracingSpec.scala
@@ -17,8 +17,8 @@ package kamon.spray
 
 import _root_.spray.httpx.RequestBuilding
 import _root_.spray.routing.SimpleRoutingApp
-import akka.testkit.TestKit
-import akka.actor.ActorSystem
+import akka.testkit.{ TestProbe, TestKit }
+import akka.actor.{ ActorRef, ActorSystem }
 import org.scalatest.{ Matchers, WordSpecLike }
 import scala.concurrent.Await
 import scala.concurrent.duration._
@@ -28,27 +28,29 @@ import kamon.trace.{ UowTrace, Trace }
 import kamon.Kamon
 import org.scalatest.concurrent.{ PatienceConfiguration, ScalaFutures }
 import spray.http.HttpHeaders.RawHeader
-import spray.http.HttpRequest
+import spray.http.{ HttpResponse, HttpRequest }
 import spray.http.HttpHeaders.Host
+import akka.io.{ Tcp, IO }
+import spray.can.Http
+import akka.io.Tcp.Bound
 
-class ServerRequestTracingSpec extends TestKit(ActorSystem("server-request-tracing-spec")) with WordSpecLike with Matchers with RequestBuilding with ScalaFutures with PatienceConfiguration with TestServer {
+class ServerRequestTracingSpec extends TestKit(ActorSystem("spec")) with WordSpecLike with Matchers with RequestBuilding with ScalaFutures with PatienceConfiguration with TestServer {
 
   "the spray server request tracing instrumentation" should {
-    "trace a request start/finish sequence when proper TraceContext is received" in {
-      val response = send {
-        Get("/ok")
-      }
+    "reply back with a trace token header" in {
+      val (connection, server) = buildServer()
+      val client = TestProbe()
 
-      within(5 seconds) {
-        fishForNamedTrace("ok")
-      }
+      client.send(connection, Get("/"))
+      server.expectMsgType[HttpRequest]
+      server.reply(HttpResponse(entity = "ok"))
+      client.expectMsgType[HttpResponse]
+
+      fail()
 
-      whenReady(response) { rsp ⇒
-        rsp.headers should contain(RawHeader("X-Trace-Token", ""))
-      }
     }
 
-    "finish a request even if no TraceContext is received in the response" in {
+    /*    "finish a request even if no TraceContext is received in the response" in {
       send {
         Get(s"http://127.0.0.1:$port/clearcontext")
       }
@@ -66,7 +68,7 @@ class ServerRequestTracingSpec extends TestKit(ActorSystem("server-request-traci
       within(5 seconds) {
         fishForNamedTrace("accounts")
       }
-    }
+    }*/
   }
   /*
   - si no llega uow, crear una
@@ -82,27 +84,23 @@ class ServerRequestTracingSpec extends TestKit(ActorSystem("server-request-traci
 trait TestServer extends SimpleRoutingApp {
   self: TestKit ⇒
 
-  Kamon(Trace).api.tell(Trace.Register, 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
-
-  val send = includeHost("127.0.0.1", port) ~> sendReceive(system, system.dispatcher, timeout)
-
-  def includeHost(host: String, port: Int) = { request: HttpRequest ⇒
-    request.withEffectiveUri(port == 443, Host(host, port))
+  def buildServer(): (ActorRef, TestProbe) = {
+    val serverHandler = TestProbe()
+    IO(Http).tell(Http.Bind(listener = serverHandler.ref, interface = "127.0.0.1", port = 0), serverHandler.ref)
+    val bound = serverHandler.expectMsgType[Bound]
+
+    val client = buildClient(bound)
+    serverHandler.expectMsgType[Http.Connected]
+    serverHandler.reply(Http.Register(serverHandler.ref))
+
+    (client, serverHandler)
+  }
+
+  def buildClient(connectionInfo: Http.Bound): ActorRef = {
+    val probe = TestProbe()
+    probe.send(IO(Http), Http.Connect(connectionInfo.localAddress.getHostName, connectionInfo.localAddress.getPort))
+    probe.expectMsgType[Http.Connected]
+    probe.sender
   }
 
 }