1 files changed, 13 insertions, 12 deletions
diff --git a/kamon-play/src/test/scala/kamon/play/LoggerLikeInstrumentationSpec.scala b/kamon-play/src/test/scala/kamon/play/LoggerLikeInstrumentationSpec.scala
index c41f7004..04ae9729 100644
--- a/kamon-play/src/test/scala/kamon/play/LoggerLikeInstrumentationSpec.scala
+++ b/kamon-play/src/test/scala/kamon/play/LoggerLikeInstrumentationSpec.scala
@@ -16,10 +16,11 @@
 package kamon.play
 
 import ch.qos.logback.classic.spi.ILoggingEvent
-import ch.qos.logback.classic.{ AsyncAppender, LoggerContext }
+import ch.qos.logback.classic.{AsyncAppender, LoggerContext}
 import ch.qos.logback.core.read.ListAppender
 import ch.qos.logback.core.status.NopStatusListener
 import kamon.trace.TraceLocal
+import kamon.trace.TraceLocal.AvailableToMdc
 import org.scalatest.BeforeAndAfter
 import org.scalatestplus.play._
 import org.slf4j
@@ -28,8 +29,9 @@ import play.api.mvc.Results.Ok
 import play.api.mvc._
 import play.api.test.Helpers._
 import play.api.test._
+import scala.concurrent.duration._
 
-import scala.concurrent.Future
+import scala.concurrent.{Await, Future}
 
 class LoggerLikeInstrumentationSpec extends PlaySpec with OneServerPerSuite with BeforeAndAfter {
 
@@ -41,11 +43,8 @@ class LoggerLikeInstrumentationSpec extends PlaySpec with OneServerPerSuite with
   val headerValue = "My header value"
   val otherValue = "My other value"
 
-  case class LocalStorageValue(header: String, other: String)
-
-  object TraceLocalKey extends TraceLocal.TraceLocalKey {
-    type ValueType = LocalStorageValue
-  }
+  val TraceLocalHeaderKey = AvailableToMdc("header")
+  val TraceLocalOtherKey = AvailableToMdc("other")
 
   before {
     LoggingHandler.startLogging()
@@ -60,7 +59,8 @@ class LoggerLikeInstrumentationSpec extends PlaySpec with OneServerPerSuite with
     case ("GET", "/logging") ⇒
       Action.async {
         Future {
-          TraceLocal.store(TraceLocalKey)(LocalStorageValue(headerValue, otherValue))
+          TraceLocal.store(TraceLocalHeaderKey)(headerValue)
+          TraceLocal.store(TraceLocalOtherKey)(otherValue)
           LoggingHandler.info(infoMessage)
           Ok("OK")
         }(executor)
@@ -68,12 +68,13 @@ class LoggerLikeInstrumentationSpec extends PlaySpec with OneServerPerSuite with
   })
 
   "the LoggerLike instrumentation" should {
-    "be put the properties of TraceLocal into the MDC as key -> value in a request" in {
+    "allow retrieve a value from the MDC when was created a key of type AvailableToMdc in the current request" in {
       LoggingHandler.appenderStart()
 
-      val Some(result) = route(FakeRequest(GET, "/logging"))
-      Thread.sleep(500) // wait to complete the future
-      TraceLocal.retrieve(TraceLocalKey) must be(Some(LocalStorageValue(headerValue, otherValue)))
+      Await.result(route(FakeRequest(GET, "/logging")).get, 500 millis)
+
+      TraceLocal.retrieve(TraceLocalHeaderKey) must be(Some(headerValue))
+      TraceLocal.retrieve(TraceLocalOtherKey) must be(Some(otherValue))
 
       LoggingHandler.appenderStop()