Included test for cases in which TraceContext are not available at creation time

2 files changed, 18 insertions, 22 deletions

diff --git a/src/main/scala/kamon/TraceContext.scala b/src/main/scala/kamon/TraceContext.scala
index 5787167b..e3582c60 100644
--- a/src/main/scala/kamon/TraceContext.scala
+++ b/src/main/scala/kamon/TraceContext.scala
@@ -14,9 +14,9 @@ object TraceContext {
     override def initialValue(): Option[TraceContext] = None
   }
 
-  def current = context.get
+  def current = context.get()
 
-  def clear = context.remove
+  def clear = context.remove()
 
   def set(ctx: TraceContext) = context.set(Some(ctx))
 
@@ -27,16 +27,3 @@ trait TraceEntry
 case class MessageExecutionTime(actorPath: ActorPath, initiated: Long, ended: Long)
 
 case class CodeBlockExecutionTime(blockName: String, begin: Long, end: Long) extends TraceEntry
-
-
-
-
-trait TraceSupport {
-  def withContext[Out](func: => Any => Out, ctx: TraceContext) = {
-    TraceContext.set(ctx)
-    val result = func
-    TraceContext.clear
-
-    result
-  }
-}
diff --git a/src/test/scala/kamon/instrumentation/FutureInstrumentationSpec.scala b/src/test/scala/kamon/instrumentation/FutureInstrumentationSpec.scala
index 6f2a3678..713d1248 100644
--- a/src/test/scala/kamon/instrumentation/FutureInstrumentationSpec.scala
+++ b/src/test/scala/kamon/instrumentation/FutureInstrumentationSpec.scala
@@ -20,8 +20,7 @@ class FutureInstrumentationSpec extends WordSpec with MustMatchers with ScalaFut
             whenReady(futureWithContext) { result =>
               result.value must be === testContext
             }
-          }
-        }
+        }}
 
         "should be available during the execution of onComplete callbacks" in { new FutureWithContext {
             val onCompleteContext = Promise[TraceContext]()
@@ -33,19 +32,29 @@ class FutureInstrumentationSpec extends WordSpec with MustMatchers with ScalaFut
             whenReady(onCompleteContext.future) { result =>
               result must be === testContext
             }
-          }
-        }
+        }}
       }
     }
 
     "created in a thread that doest have a TraceContext" must {
-      "not capture any TraceContext" in { new FutureWithoutContext{
+      "not capture any TraceContext for the body execution" in { new FutureWithoutContext{
 
           whenReady(futureWithoutContext) { result =>
             result must be === None
           }
+      }}
+
+      "not make any TraceContext available during the onComplete callback" in { new FutureWithoutContext {
+        val onCompleteContext = Promise[Option[TraceContext]]()
+
+        futureWithoutContext.onComplete({
+          case _ => onCompleteContext.complete(Success(TraceContext.current))
+        })
+
+        whenReady(onCompleteContext.future) { result =>
+          result must be === None
         }
-      }
+      }}
     }
   }
 
@@ -60,7 +69,7 @@ class FutureInstrumentationSpec extends WordSpec with MustMatchers with ScalaFut
   }
 
   trait FutureWithoutContext {
-    TraceContext.clear
+    TraceContext.clear // Make sure no TraceContext is available
     val futureWithoutContext = Future { TraceContext.current }
   }
 }