2 files changed, 83 insertions, 23 deletions

diff --git a/src/test/scala/GoogleTracerSpec.scala b/src/test/scala/GoogleTracerSpec.scala
new file mode 100644
index 0000000..fcdd73b
--- /dev/null
+++ b/src/test/scala/GoogleTracerSpec.scala
@@ -0,0 +1,51 @@
+package xyz.driver.tracing
+
+import java.nio.file._
+
+import akka.actor._
+import akka.stream._
+import akka.http.scaladsl.server.Directives._
+import akka.http.scaladsl.server._
+import akka.http.scaladsl.testkit.ScalatestRouteTest
+import org.scalatest._
+import xyz.driver.tracing.TracingDirectives._
+
+import scala.concurrent._
+import scala.concurrent.duration._
+
+class GoogleTracerSpec extends FlatSpec {
+
+  implicit val system = ActorSystem(this.getClass.getSimpleName)
+  import system.dispatcher
+  implicit val materializer = ActorMaterializer()
+
+  "Google Tracer" should "submit" in {
+    val tracer = new GoogleTracer(
+      "driverinc-sandbox",
+      Paths.get(
+        system.settings.config.getString("tracing.google.service-account-file"))
+    )
+
+    // create 100 traces, each with 10 nested spans
+    for (_ <- 0 until 100) {
+      (0 until 10).foldLeft(Span(name = "trace-test.driver.xyz")) {
+        case (previous, i) =>
+          val p: Span = previous
+          val span = Span(
+            name = s"trace-test.driver.xyz-$i",
+            traceId = previous.traceId,
+            parentSpanId = Some(previous.spanId)
+          )
+          Thread.sleep(2)
+          val done = span.end()
+          tracer.submit(done)
+          done
+      }
+    }
+
+    // TODO: automatically verify that traces were created on google stack driver
+
+    Await.ready(tracer.close(), 30.seconds)
+  }
+
+}
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)
   }