Add type hints and parent assertion

dd-java-agent/instrumentation/rabbitmq-amqp-2.7/src/test/groovy/RabbitMQTest.groovy

@@ -25,6 +25,8 @@ import spock.lang.Shared
 
 import java.util.concurrent.Phaser
 
+import static datadog.trace.agent.test.TestUtils.runUnderTrace
+
 // Do not run tests locally on Java7 since testcontainers are not compatible with Java7
 // It is fine to run on CI because CI provides rabbitmq externally, not through testcontainers
 @Requires({ "true" == System.getenv("CI") || jvm.java8Compatible })
@@ -83,33 +85,33 @@ class RabbitMQTest extends AgentTestRunner {
 
   def "test rabbit publish/get"() {
     setup:
+    GetResponse response = runUnderTrace("parent") {
       channel.exchangeDeclare(exchangeName, "direct", false)
       String queueName = channel.queueDeclare().getQueue()
       channel.queueBind(queueName, exchangeName, routingKey)
 
       channel.basicPublish(exchangeName, routingKey, null, "Hello, world!".getBytes())
 
-    GetResponse response = channel.basicGet(queueName, true)
+      return channel.basicGet(queueName, true)
+    }
 
     expect:
     new String(response.getBody()) == "Hello, world!"
 
     and:
-    assertTraces(5) {
+    assertTraces(2) {
       trace(0, 1) {
-        rabbitSpan(it, "exchange.declare")
+        rabbitSpan(it, "basic.get <generated>", TEST_WRITER[1][1])
       }
-      trace(1, 1) {
+      trace(1, 5) {
-        rabbitSpan(it, "queue.declare")
+        span(0) {
+          operationName "parent"
         }
-      trace(2, 1) {
+        // reverse order
-        rabbitSpan(it, "queue.bind")
+        rabbitSpan(it, 1, "basic.publish $exchangeName", span(0))
-      }
+        rabbitSpan(it, 2, "queue.bind", span(0))
-      trace(3, 1) {
+        rabbitSpan(it, 3, "queue.declare", span(0))
-        rabbitSpan(it, "basic.publish $exchangeName")
+        rabbitSpan(it, 4, "exchange.declare", span(0))
-      }
-      trace(4, 1) {
-        rabbitSpan(it, "basic.get <generated>", TEST_WRITER[3][0])
       }
     }
 
@@ -264,7 +266,11 @@ class RabbitMQTest extends AgentTestRunner {
   }
 
   def rabbitSpan(TraceAssert trace, String resource, DDSpan parentSpan = null, Throwable exception = null, String errorMsg = null) {
-    trace.span(0) {
+    rabbitSpan(trace, 0, resource, parentSpan, exception, errorMsg)
+  }
+
+  def rabbitSpan(TraceAssert trace, int index, String resource, DDSpan parentSpan = null, Throwable exception = null, String errorMsg = null) {
+    trace.span(index) {
       serviceName "rabbitmq"
       operationName "amqp.command"
       resourceName resource

dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/AgentTestRunner.java

@@ -13,6 +13,8 @@ import datadog.trace.common.writer.ListWriter;
 import datadog.trace.common.writer.Writer;
 import groovy.lang.Closure;
 import groovy.lang.DelegatesTo;
+import groovy.transform.stc.ClosureParams;
+import groovy.transform.stc.SimpleType;
 import io.opentracing.Tracer;
 import java.lang.instrument.ClassFileTransformer;
 import java.lang.instrument.Instrumentation;
@@ -160,6 +162,9 @@ public abstract class AgentTestRunner extends Specification {
 
   public static void assertTraces(
       final int size,
+      @ClosureParams(
+              value = SimpleType.class,
+              options = "datadog.trace.agent.test.asserts.ListWriterAssert")
           @DelegatesTo(value = ListWriterAssert.class, strategy = Closure.DELEGATE_FIRST)
           final Closure spec) {
     ListWriterAssert.assertTraces(TEST_WRITER, size, spec);

dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/ListWriterAssert.groovy

@@ -2,6 +2,8 @@ package datadog.trace.agent.test.asserts
 
 import datadog.opentracing.DDSpan
 import datadog.trace.common.writer.ListWriter
+import groovy.transform.stc.ClosureParams
+import groovy.transform.stc.SimpleType
 import org.codehaus.groovy.runtime.powerassert.PowerAssertionError
 import org.spockframework.runtime.Condition
 import org.spockframework.runtime.ConditionNotSatisfiedError
@@ -20,6 +22,7 @@ class ListWriterAssert {
   }
 
   static void assertTraces(ListWriter writer, int expectedSize,
+                           @ClosureParams(value = SimpleType, options = ['datadog.trace.agent.test.asserts.ListWriterAssert'])
                            @DelegatesTo(value = ListWriterAssert, strategy = Closure.DELEGATE_FIRST) Closure spec) {
     try {
       writer.waitForTraces(expectedSize)
@@ -55,6 +58,7 @@ class ListWriterAssert {
   }
 
   void trace(int index, int expectedSize,
+             @ClosureParams(value = SimpleType, options = ['datadog.trace.agent.test.asserts.TraceAssert'])
              @DelegatesTo(value = TraceAssert, strategy = Closure.DELEGATE_FIRST) Closure spec) {
     if (index >= size) {
       throw new ArrayIndexOutOfBoundsException(index)

dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/SpanAssert.groovy

@@ -1,6 +1,8 @@
 package datadog.trace.agent.test.asserts
 
 import datadog.opentracing.DDSpan
+import groovy.transform.stc.ClosureParams
+import groovy.transform.stc.SimpleType
 
 import static TagsAssert.assertTags
 
@@ -12,6 +14,7 @@ class SpanAssert {
   }
 
   static void assertSpan(DDSpan span,
+                         @ClosureParams(value = SimpleType, options = ['datadog.trace.agent.test.asserts.SpanAssert'])
                          @DelegatesTo(value = SpanAssert, strategy = Closure.DELEGATE_FIRST) Closure spec) {
     def asserter = new SpanAssert(span)
     def clone = (Closure) spec.clone()
@@ -72,7 +75,8 @@ class SpanAssert {
     assert span.isError() == errored
   }
 
-  void tags(@DelegatesTo(value = TagsAssert, strategy = Closure.DELEGATE_FIRST) Closure spec) {
+  void tags(@ClosureParams(value = SimpleType, options = ['datadog.trace.agent.test.asserts.TagsAssert'])
+            @DelegatesTo(value = TagsAssert, strategy = Closure.DELEGATE_FIRST) Closure spec) {
     assertTags(span, spec)
   }
 }

dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/TagsAssert.groovy

@@ -1,6 +1,8 @@
 package datadog.trace.agent.test.asserts
 
 import datadog.opentracing.DDSpan
+import groovy.transform.stc.ClosureParams
+import groovy.transform.stc.SimpleType
 
 class TagsAssert {
   private final Map<String, Object> tags
@@ -11,6 +13,7 @@ class TagsAssert {
   }
 
   static void assertTags(DDSpan span,
+                         @ClosureParams(value = SimpleType, options = ['datadog.trace.agent.test.asserts.TagsAssert'])
                          @DelegatesTo(value = TagsAssert, strategy = Closure.DELEGATE_FIRST) Closure spec) {
     def asserter = new TagsAssert(span)
     def clone = (Closure) spec.clone()

dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/TraceAssert.groovy

@@ -1,6 +1,8 @@
 package datadog.trace.agent.test.asserts
 
 import datadog.opentracing.DDSpan
+import groovy.transform.stc.ClosureParams
+import groovy.transform.stc.SimpleType
 
 import static SpanAssert.assertSpan
 
@@ -15,6 +17,7 @@ class TraceAssert {
   }
 
   static void assertTrace(List<DDSpan> trace, int expectedSize,
+                          @ClosureParams(value = SimpleType, options = ['datadog.trace.agent.test.asserts.TraceAssert'])
                           @DelegatesTo(value = File, strategy = Closure.DELEGATE_FIRST) Closure spec) {
     assert trace.size() == expectedSize
     def asserter = new TraceAssert(trace)
@@ -29,7 +32,7 @@ class TraceAssert {
     trace.get(index)
   }
 
-  void span(int index, @DelegatesTo(value = SpanAssert, strategy = Closure.DELEGATE_FIRST) Closure spec) {
+  void span(int index, @ClosureParams(value = SimpleType, options = ['datadog.trace.agent.test.asserts.SpanAssert']) @DelegatesTo(value = SpanAssert, strategy = Closure.DELEGATE_FIRST) Closure spec) {
     if (index >= size) {
       throw new ArrayIndexOutOfBoundsException(index)
     }