Added a simple test of futures instrumentation

1 files changed, 45 insertions, 13 deletions

diff --git a/src/test/scala/kamon/instrumentation/FutureInstrumentationSpec.scala b/src/test/scala/kamon/instrumentation/FutureInstrumentationSpec.scala
index 97e81b44..6f2a3678 100644
--- a/src/test/scala/kamon/instrumentation/FutureInstrumentationSpec.scala
+++ b/src/test/scala/kamon/instrumentation/FutureInstrumentationSpec.scala
@@ -1,35 +1,67 @@
 package kamon.instrumentation
 
-import scala.concurrent.Future
+import scala.concurrent.{Promise, Future}
 import scala.concurrent.ExecutionContext.Implicits.global
-import org.scalatest.WordSpec
+import org.scalatest.{OptionValues, WordSpec}
 import org.scalatest.matchers.MustMatchers
 import org.scalatest.concurrent.PatienceConfiguration
 import kamon.TraceContext
 import java.util.UUID
+import scala.util.Success
 
 
-class FutureInstrumentationSpec extends WordSpec with MustMatchers with ScalaFutures with PatienceConfiguration {
+class FutureInstrumentationSpec extends WordSpec with MustMatchers with ScalaFutures with PatienceConfiguration with OptionValues {
 
-  "a instrumented Future" should {
-    "preserve the transaction context available during the future creation" in {
-      new ContextAwareTest {
-        val future = Future { TraceContext.current.get }
+  "a instrumented Future" when {
+    "created in a thread that does have a TraceContext" must {
+      "preserve the TraceContext" which {
+        "should be available during the body's execution" in { new FutureWithContext {
 
-        whenReady(future) { result =>
-          result must be === context
+            whenReady(futureWithContext) { result =>
+              result.value must be === testContext
+            }
+          }
+        }
+
+        "should be available during the execution of onComplete callbacks" in { new FutureWithContext {
+            val onCompleteContext = Promise[TraceContext]()
+
+            futureWithContext.onComplete({
+              case _ => onCompleteContext.complete(Success(TraceContext.current.get))
+            })
+
+            whenReady(onCompleteContext.future) { result =>
+              result must be === testContext
+            }
+          }
         }
       }
     }
 
-    "use the same context available at creation when executing the onComplete callback" in {
+    "created in a thread that doest have a TraceContext" must {
+      "not capture any TraceContext" in { new FutureWithoutContext{
 
+          whenReady(futureWithoutContext) { result =>
+            result must be === None
+          }
+        }
+      }
     }
   }
 
-  trait ContextAwareTest {
-    val context = TraceContext(UUID.randomUUID(), Nil)
-    TraceContext.set(context)
+
+
+
+  trait FutureWithContext {
+    val testContext = TraceContext(UUID.randomUUID(), Nil)
+    TraceContext.set(testContext)
+
+    val futureWithContext = Future { TraceContext.current }
+  }
+
+  trait FutureWithoutContext {
+    TraceContext.clear
+    val futureWithoutContext = Future { TraceContext.current }
   }
 }