1 files changed, 12 insertions, 14 deletions
diff --git a/kamon-core/src/test/scala/kamon/trace/instrumentation/AskPatternTracingSpec.scala b/kamon-core/src/test/scala/kamon/trace/instrumentation/AskPatternTracingSpec.scala
index 9df67391..0387386c 100644
--- a/kamon-core/src/test/scala/kamon/trace/instrumentation/AskPatternTracingSpec.scala
+++ b/kamon-core/src/test/scala/kamon/trace/instrumentation/AskPatternTracingSpec.scala
@@ -22,32 +22,30 @@ import akka.event.Logging.Warning
 import scala.concurrent.duration._
 import akka.pattern.ask
 import akka.util.Timeout
-import kamon.trace.{ Trace, ContextAware }
+import kamon.trace.{TraceContextAware, TraceRecorder}
 import org.scalatest.OptionValues._
 
 class AskPatternTracingSpec extends TestKit(ActorSystem("ask-pattern-tracing-spec")) with WordSpecLike with Matchers {
 
   "the AskPatternTracing" should {
-    "log a warning with a stack trace and TraceContext taken from the moment the ask was triggered" in new TraceContextFixture {
+    "log a warning with a stack trace and TraceContext taken from the moment the ask was triggered" in {
       implicit val ec = system.dispatcher
       implicit val timeout = Timeout(10 milliseconds)
       val noReply = system.actorOf(Props[NoReply])
       system.eventStream.subscribe(testActor, classOf[Warning])
 
-      within(500 milliseconds) {
-        val initialCtx = Trace.withContext(testTraceContext) {
-          noReply ? "hello"
-          Trace.context()
-        }
-
-        val warn = expectMsgPF() {
-          case warn: Warning if warn.message.toString.contains("Timeout triggered for ask pattern") ⇒ warn
-        }
-        val capturedCtx = warn.asInstanceOf[ContextAware].traceContext
+      val testTraceContext = TraceRecorder.withNewTraceContext("ask-timeout-warning") {
+        noReply ? "hello"
+        TraceRecorder.currentContext
+      }
 
-        capturedCtx should be('defined)
-        capturedCtx should equal(initialCtx)
+      val warn = expectMsgPF() {
+        case warn: Warning if warn.message.toString.contains("Timeout triggered for ask pattern") ⇒ warn
       }
+      val capturedCtx = warn.asInstanceOf[TraceContextAware].traceContext
+
+      capturedCtx should be('defined)
+      capturedCtx should equal(testTraceContext)
     }
   }
 }