tests for Span building and the ExtendedB3 codec, related bugfixes

2 files changed, 69 insertions, 3 deletions

diff --git a/kamon-core/src/test/scala/kamon/testkit/SpanBuilding.scala b/kamon-core/src/test/scala/kamon/testkit/SpanBuilding.scala
new file mode 100644
index 00000000..9b845ac9
--- /dev/null
+++ b/kamon-core/src/test/scala/kamon/testkit/SpanBuilding.scala
@@ -0,0 +1,20 @@
+package kamon.testkit
+
+import kamon.trace.SpanContext.{SamplingDecision, Source}
+import kamon.trace.{IdentityProvider, SpanContext, SpanContextCodec}
+
+trait SpanBuilding {
+  private val identityProvider = IdentityProvider.Default()
+  private val extendedB3Codec = SpanContextCodec.ExtendedB3(identityProvider)
+
+  def createSpanContext(samplingDecision: SamplingDecision = SamplingDecision.Sample): SpanContext =
+    SpanContext(
+      traceID = identityProvider.traceIdentifierGenerator().generate(),
+      spanID = identityProvider.spanIdentifierGenerator().generate(),
+      parentID = identityProvider.spanIdentifierGenerator().generate(),
+      samplingDecision = samplingDecision,
+      baggage = SpanContext.Baggage(),
+      source = Source.Local
+    )
+
+}
diff --git a/kamon-core/src/test/scala/kamon/testkit/SpanInspector.scala b/kamon-core/src/test/scala/kamon/testkit/SpanInspector.scala
index b0bb3e39..3ef1012b 100644
--- a/kamon-core/src/test/scala/kamon/testkit/SpanInspector.scala
+++ b/kamon-core/src/test/scala/kamon/testkit/SpanInspector.scala
@@ -1,15 +1,61 @@
 package kamon.testkit
 
-import kamon.trace.Span
+import kamon.trace.{ActiveSpan, Span, SpanContext}
 import kamon.trace.Span.FinishedSpan
+import kamon.util.Clock
+
+import scala.reflect.ClassTag
+import scala.util.Try
 
 class SpanInspector(span: Span) {
+  private val (realSpan, spanData) = {
+    val realSpan = span match {
+      case _: Span.Real => span
+      case a: ActiveSpan =>
+        getField[ActiveSpan.Default, Span](a, "wrappedSpan")
+    }
+
+    val spanData = invoke[Span.Real, FinishedSpan](realSpan, "toFinishedSpan", classOf[Long] -> Long.box(Clock.microTimestamp()))
+    (realSpan, spanData)
+  }
+
+  def nonEmpty: Boolean =
+    !span.isInstanceOf[Span.Empty]
+
+  def spanTag(key: String): Option[Span.TagValue] =
+    spanData.tags.get(key)
+
+  def spanTags(): Map[String, Span.TagValue] =
+    spanData.tags
+
+  def metricTags(): Map[String, String] =
+    getField[Span.Real, Map[String, String]](realSpan, "customMetricTags")
 
+  def startTimestamp(): Long =
+    getField[Span.Real, Long](realSpan, "startTimestampMicros")
 
-  private def getSpanData(): Option[FinishedSpan] = {
-    
+  def context(): SpanContext =
+    spanData.context
+
+  def operationName(): String =
+    spanData.operationName
+
+
+
+
+  private def getField[T, R](target: Any, fieldName: String)(implicit classTag: ClassTag[T]): R = {
+    val toFinishedSpanMethod = classTag.runtimeClass.getDeclaredField(fieldName)
+    toFinishedSpanMethod.setAccessible(true)
+    toFinishedSpanMethod.get(target).asInstanceOf[R]
   }
 
+  private def invoke[T, R](target: Any, fieldName: String, parameters: (Class[_], AnyRef)*)(implicit classTag: ClassTag[T]): R = {
+    val parameterClasses = parameters.map(_._1)
+    val parameterInstances = parameters.map(_._2)
+    val toFinishedSpanMethod = classTag.runtimeClass.getDeclaredMethod(fieldName, parameterClasses: _*)
+    toFinishedSpanMethod.setAccessible(true)
+    toFinishedSpanMethod.invoke(target, parameterInstances: _*).asInstanceOf[R]
+  }
 }
 
 object SpanInspector {