+ core: refactor MDC facilities and closes #100

2 files changed, 57 insertions, 5 deletions

diff --git a/kamon-core/src/test/scala/kamon/instrumentation/akka/ActorLoggingInstrumentationSpec.scala b/kamon-core/src/test/scala/kamon/instrumentation/akka/ActorLoggingInstrumentationSpec.scala
index 3dab44bc..cef335e0 100644
--- a/kamon-core/src/test/scala/kamon/instrumentation/akka/ActorLoggingInstrumentationSpec.scala
+++ b/kamon-core/src/test/scala/kamon/instrumentation/akka/ActorLoggingInstrumentationSpec.scala
@@ -18,11 +18,14 @@ package kamon.instrumentation.akka
 import akka.actor.{ Actor, ActorLogging, ActorSystem, Props }
 import akka.event.Logging.LogEvent
 import akka.testkit.TestKit
-import kamon.trace.{ TraceContextAware, TraceRecorder }
+import kamon.trace.TraceLocal.AvailableToMdc
+import kamon.trace.logging.MdcKeysSupport
+import kamon.trace.{TraceLocal, TraceContextAware, TraceRecorder}
 import org.scalatest.{ Inspectors, Matchers, WordSpecLike }
+import org.slf4j.MDC
 
 class ActorLoggingInstrumentationSpec extends TestKit(ActorSystem("actor-logging-instrumentation-spec")) with WordSpecLike
-    with Matchers with Inspectors {
+    with Matchers with Inspectors with MdcKeysSupport {
 
   "the ActorLogging instrumentation" should {
     "attach the TraceContext (if available) to log events" in {
@@ -42,6 +45,24 @@ class ActorLoggingInstrumentationSpec extends TestKit(ActorSystem("actor-logging
         case event: LogEvent ⇒ false
       }
     }
+
+    "allow retrieve a value from the MDC when was created a key of type AvailableToMdc" in {
+      val testString = "Hello World"
+      val SampleTraceLocalKeyAvailableToMDC = AvailableToMdc("some-cool-key")
+
+      val loggerActor = system.actorOf(Props[LoggerActor])
+      system.eventStream.subscribe(testActor, classOf[LogEvent])
+
+      TraceRecorder.withNewTraceContext("logging-with-mdc") {
+        TraceLocal.store(SampleTraceLocalKeyAvailableToMDC)(testString)
+
+        loggerActor ! "info"
+
+        withMdc {
+          MDC.get("some-cool-key") should equal(testString)
+        }
+      }
+    }
   }
 }
 
diff --git a/kamon-core/src/test/scala/kamon/trace/TraceLocalSpec.scala b/kamon-core/src/test/scala/kamon/trace/TraceLocalSpec.scala
index 927573c2..b23c1c03 100644
--- a/kamon-core/src/test/scala/kamon/trace/TraceLocalSpec.scala
+++ b/kamon-core/src/test/scala/kamon/trace/TraceLocalSpec.scala
@@ -16,13 +16,18 @@
 
 package kamon.trace
 
-import akka.testkit.TestKit
 import akka.actor.ActorSystem
-import org.scalatest.{ OptionValues, Matchers, WordSpecLike }
+import akka.testkit.TestKit
+import kamon.trace.TraceLocal.AvailableToMdc
+import kamon.trace.logging.MdcKeysSupport
 import org.scalatest.concurrent.PatienceConfiguration
+import org.scalatest.{Matchers, OptionValues, WordSpecLike}
+import org.slf4j.MDC
 
 class TraceLocalSpec extends TestKit(ActorSystem("trace-local-spec")) with WordSpecLike with Matchers
-    with PatienceConfiguration with OptionValues {
+    with PatienceConfiguration with OptionValues with MdcKeysSupport {
+
+  val SampleTraceLocalKeyAvailableToMDC = AvailableToMdc("someKey")
 
   object SampleTraceLocalKey extends TraceLocal.TraceLocalKey { type ValueType = String }
 
@@ -61,5 +66,31 @@ class TraceLocalSpec extends TestKit(ActorSystem("trace-local-spec")) with WordS
         TraceLocal.retrieve(SampleTraceLocalKey).value should equal(testString)
       }
     }
+
+    "allow retrieve a value from the MDC when was created a key with AvailableToMdc(cool-key)" in {
+      TraceRecorder.withNewTraceContext("store-and-retrieve-trace-local-and-copy-to-mdc") {
+        val testString = "Hello MDC"
+
+        TraceLocal.store(SampleTraceLocalKeyAvailableToMDC)(testString)
+        TraceLocal.retrieve(SampleTraceLocalKeyAvailableToMDC).value should equal(testString)
+
+        withMdc {
+          MDC.get("someKey") should equal(testString)
+        }
+      }
+    }
+
+    "allow retrieve a value from the MDC when was created a key with AvailableToMdc.storeForMdc(String, String)" in {
+      TraceRecorder.withNewTraceContext("store-and-retrieve-trace-local-and-copy-to-mdc") {
+        val testString = "Hello MDC"
+
+        TraceLocal.storeForMdc("someKey", testString)
+        TraceLocal.retrieve(SampleTraceLocalKeyAvailableToMDC).value should equal(testString)
+
+        withMdc {
+          MDC.get("someKey") should equal(testString)
+        }
+      }
+    }
   }
 }