1 files changed, 32 insertions, 23 deletions
diff --git a/src/test/scala/TracingDirectivesSpec.scala b/src/test/scala/TracingDirectivesSpec.scala
index 783f59b..61206c8 100644
--- a/src/test/scala/TracingDirectivesSpec.scala
+++ b/src/test/scala/TracingDirectivesSpec.scala
@@ -16,40 +16,49 @@ class TracingDirectivesSpec
     with BeforeAndAfterAll
     with ScalatestRouteTest {
 
-  def route(tracer: Tracer): Route = trace(tracer, "example.org") {
-    pathPrefix("1") {
-      trace(tracer, "test-sub-trace-1") {
+  class DummyTracer extends Tracer {
+    private var _spans: List[Span] = Nil
+    def spans = synchronized { _spans }
+
+    override def submit(span: Span): Unit = synchronized {
+      _spans = span :: _spans
+    }
+  }
+
+  def route(tracer: Tracer): Route = trace(tracer, Some("example.org")) {
+    trace(tracer, Some("trace-1-1")) {
+      Thread.sleep(2)
+      trace(tracer, Some("trace-1-2")) {
         Thread.sleep(2)
-        trace(tracer, "test-subsub-trace-1") {
-          Thread.sleep(2)
-          trace(tracer, "test-subsubsub-trace-1") {
-            Thread.sleep(10)
-            complete("ok")
-          }
-        }
-      }
-    } ~
-      pathPrefix("2") {
-        trace(tracer, "test-sub-trace-2") {
-          Thread.sleep(20)
+        trace(tracer, Some("trace-1-3")) {
+          Thread.sleep(10)
           complete("ok")
         }
       }
+    }
   }
 
-  "Google Tracer" should "submit" in {
-    val tracer = new GoogleTracer(
-      "driverinc-sandbox",
-      Paths.get(
-        system.settings.config.getString("tracing.google.service-account-file"))
-    )
+  "Tracing directives" should "nest spans correctly" in {
+    val tracer = new DummyTracer
 
-    val futures: Seq[Assertion] = for (i <- 0 until 100) yield {
-      Get(s"https://example.org/${i % 2 + 1}") ~> route(tracer) ~> check {
+    for (i <- 0 until 100) yield {
+      Get(s"https://example.org") ~> route(tracer) ~> check {
         assert(responseAs[String] == "ok")
       }
     }
 
+    val traces = tracer.spans.groupBy(_.traceId)
+    assert(traces.toSeq.length == 100)
+
+    traces.foreach {
+      case (traceId, spans) =>
+        def getSpan(name: String) = spans.find(_.name == name).get
+        assert(
+          getSpan("trace-1-3").parentSpanId.get === getSpan("trace-1-2").spanId)
+        assert(
+          getSpan("trace-1-2").parentSpanId.get === getSpan("trace-1-1").spanId)
+    }
+
     Await.ready(tracer.close(), 30.seconds)
   }