1 files changed, 85 insertions, 0 deletions
diff --git a/kamon-trace/src/test/scala/kamon/RunnableInstrumentationSpec.scala b/kamon-trace/src/test/scala/kamon/RunnableInstrumentationSpec.scala
new file mode 100644
index 00000000..f968fa83
--- /dev/null
+++ b/kamon-trace/src/test/scala/kamon/RunnableInstrumentationSpec.scala
@@ -0,0 +1,85 @@
+package kamon
+
+import scala.concurrent.{Await, Promise, Future}
+import org.scalatest.{Matchers, OptionValues, WordSpec}
+import org.scalatest.concurrent.{ScalaFutures, PatienceConfiguration}
+import java.util.UUID
+import scala.util.Success
+import scala.concurrent.duration._
+import java.util.concurrent.TimeUnit
+import akka.actor.{Actor, ActorSystem}
+import kamon.trace.{Trace, TraceContext}
+
+
+class RunnableInstrumentationSpec extends WordSpec with Matchers with ScalaFutures with PatienceConfiguration with OptionValues {
+
+  "a instrumented runnable" when {
+    "created in a thread that does have a TraceContext" must {
+      "preserve the TraceContext" which {
+        "should be available during the run method execution" in new FutureWithContextFixture {
+
+/*            whenReady(futureWithContext) { result =>
+              result.value should equal(testContext)
+            }*/
+        }
+
+        "should be available during the execution of onComplete callbacks" in new FutureWithContextFixture {
+
+          val onCompleteContext = Promise[Option[TraceContext]]()
+
+/*          Tracer.traceContext.withValue(Some(testContext)) {
+            futureWithContext.onComplete({
+              case _ => println("Completing second promise from: "+Thread.currentThread().getName + " With Context: " + Tracer.traceContext.value); onCompleteContext.complete(Success(Tracer.traceContext.value))
+            })
+          }*/
+
+          whenReady(onCompleteContext.future) { result =>
+            result should equal(Some(testContext))
+          }
+        }
+      }
+    }
+
+    "created in a thread that doest have a TraceContext" must {
+      "not capture any TraceContext for the body execution" in new FutureWithoutContextFixture{
+          whenReady(futureWithoutContext) { result =>
+            result should equal(None)
+          }
+      }
+
+      "not make any TraceContext available during the onComplete callback" in new FutureWithoutContextFixture {
+        val onCompleteContext = Promise[Option[TraceContext]]()
+
+        futureWithoutContext.onComplete {
+          case _ => onCompleteContext.complete(Success(Trace.context()))
+        }
+
+        whenReady(onCompleteContext.future) { result =>
+          result should equal(None)
+        }
+      }
+    }
+  }
+
+
+  /**
+   *  We are using Futures for the test since they exercise Runnables in the back and also resemble the real use case we have.
+   */
+  implicit val testActorSystem = ActorSystem("test-actorsystem")
+  implicit val execContext = testActorSystem.dispatcher
+
+  class FutureWithContextFixture {
+    val testContext = TraceContext(Actor.noSender, 1)
+
+/*    var futureWithContext: Future[Option[TraceContext]] = _
+    Tracer.context.withValue(Some(testContext)) {
+      futureWithContext = Future { Tracer.traceContext.value }
+    }*/
+  }
+
+  trait FutureWithoutContextFixture {
+    val futureWithoutContext = Future { Trace.context.value }
+  }
+}
+
+