! all: improve the metric recorders infrastructure

2 files changed, 33 insertions, 2 deletions

diff --git a/kamon-testkit/src/main/scala/testkit/AkkaExtensionSwap.scala b/kamon-testkit/src/main/scala/testkit/AkkaExtensionSwap.scala
new file mode 100644
index 00000000..2f77df95
--- /dev/null
+++ b/kamon-testkit/src/main/scala/testkit/AkkaExtensionSwap.scala
@@ -0,0 +1,31 @@
+/*
+ * =========================================================================================
+ * Copyright © 2013-2014 the kamon project <http://kamon.io/>
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
+ * except in compliance with the License. You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the
+ * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
+ * either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ * =========================================================================================
+ */
+
+package testkit
+
+import java.util.concurrent.ConcurrentHashMap
+
+import akka.actor.{ ActorSystem, Extension, ExtensionId }
+
+object AkkaExtensionSwap {
+  def swap(system: ActorSystem, key: ExtensionId[_], value: Extension): Unit = {
+    val extensionsField = system.getClass.getDeclaredField("extensions")
+    extensionsField.setAccessible(true)
+
+    val extensions = extensionsField.get(system).asInstanceOf[ConcurrentHashMap[ExtensionId[_], AnyRef]]
+    extensions.put(key, value)
+  }
+}
diff --git a/kamon-testkit/src/main/scala/testkit/TestProbeInstrumentation.scala b/kamon-testkit/src/main/scala/testkit/TestProbeInstrumentation.scala
index 825cc718..9e736971 100644
--- a/kamon-testkit/src/main/scala/testkit/TestProbeInstrumentation.scala
+++ b/kamon-testkit/src/main/scala/testkit/TestProbeInstrumentation.scala
@@ -17,7 +17,7 @@
 package akka.testkit
 
 import org.aspectj.lang.annotation._
-import kamon.trace.{ EmptyTraceContext, TraceContextAware, TraceRecorder }
+import kamon.trace.{ EmptyTraceContext, TraceContextAware, TraceContext }
 import org.aspectj.lang.ProceedingJoinPoint
 import akka.testkit.TestActor.RealMessage
 
@@ -46,7 +46,7 @@ class TestProbeInstrumentation {
       case _                ⇒ EmptyTraceContext
     }
 
-    TraceRecorder.withTraceContext(traceContext) {
+    TraceContext.withContext(traceContext) {
       pjp.proceed
     }
   }