Fix sporadic test failure in Elasticsearch53SpringRepositoryTest

dd-java-agent/instrumentation/elasticsearch/transport-5.3/src/test/groovy/springdata/Elasticsearch53SpringRepositoryTest.groovy

@@ -1,6 +1,7 @@
 package springdata
 
 import com.anotherchrisberry.spock.extensions.retry.RetryOnFailure
+import datadog.opentracing.DDSpan
 import datadog.trace.agent.test.AgentTestRunner
 import datadog.trace.api.DDSpanTypes
 import datadog.trace.instrumentation.api.Tags
@@ -34,13 +35,8 @@ class Elasticsearch53SpringRepositoryTest extends AgentTestRunner {
         return repo
       }
 
-      TEST_WRITER.clear()
-      runUnderTrace("setup") {
-        def applicationContext = new AnnotationConfigApplicationContext(Config)
-        repo = applicationContext.getBean(DocRepository)
-      }
-      TEST_WRITER.waitForTraces(1)
-      TEST_WRITER.clear()
+      def applicationContext = new AnnotationConfigApplicationContext(Config)
+      repo = applicationContext.getBean(DocRepository)
 
       return repo
     }
@@ -68,6 +64,7 @@ class Elasticsearch53SpringRepositoryTest extends AgentTestRunner {
     !result.iterator().hasNext()
 
     and:
+    waitForTracesAndSortSpans(1)
     assertTraces(1) {
       trace(0, 2) {
         span(0) {
@@ -114,8 +111,31 @@ class Elasticsearch53SpringRepositoryTest extends AgentTestRunner {
     repo.index(doc) == doc
 
     and:
-    assertTraces(1) {
-      trace(0, 3) {
+    waitForTracesAndSortSpans(2)
+    // need to normalize trace ordering since they are finished by different threads
+    if (TEST_WRITER[1][0].resourceName == "PutMappingAction") {
+      def tmp = TEST_WRITER[1]
+      TEST_WRITER[1] = TEST_WRITER[0]
+      TEST_WRITER[0] = tmp
+    }
+    assertTraces(2) {
+      trace(0, 1) {
+        span(0) {
+          serviceName "elasticsearch"
+          resourceName "PutMappingAction"
+          operationName "elasticsearch.query"
+          spanType DDSpanTypes.ELASTICSEARCH
+          tags {
+            "$Tags.COMPONENT" "elasticsearch-java"
+            "$Tags.DB_TYPE" "elasticsearch"
+            "$Tags.SPAN_KIND" Tags.SPAN_KIND_CLIENT
+            "elasticsearch.action" "PutMappingAction"
+            "elasticsearch.request" "PutMappingRequest"
+            defaultTags()
+          }
+        }
+      }
+      trace(1, 3) {
         span(0) {
           resourceName "ElasticsearchRepository.index"
           operationName "repository.operation"
@@ -174,6 +194,7 @@ class Elasticsearch53SpringRepositoryTest extends AgentTestRunner {
     repo.findById("1").get() == doc
 
     and:
+    waitForTracesAndSortSpans(1)
     assertTraces(1) {
       trace(0, 2) {
         span(0) {
@@ -217,6 +238,7 @@ class Elasticsearch53SpringRepositoryTest extends AgentTestRunner {
     repo.findById("1").get() == doc
 
     and:
+    waitForTracesAndSortSpans(2)
     assertTraces(2) {
       trace(0, 3) {
         span(0) {
@@ -309,6 +331,7 @@ class Elasticsearch53SpringRepositoryTest extends AgentTestRunner {
     !repo.findAll().iterator().hasNext()
 
     and:
+    waitForTracesAndSortSpans(2)
     assertTraces(2) {
       trace(0, 3) {
         span(0) {
@@ -395,4 +418,17 @@ class Elasticsearch53SpringRepositoryTest extends AgentTestRunner {
     where:
     indexName = "test-index"
   }
+
+  def waitForTracesAndSortSpans(int number) {
+    TEST_WRITER.waitForTraces(number)
+    for (List<DDSpan> trace : TEST_WRITER) {
+      // need to normalize span ordering since they are finished by different threads
+      if (trace.size() > 1 && trace[1].operationName == "repository.operation") {
+        def tmp = trace[1]
+        trace[1] = trace[0]
+        trace[0] = tmp
+      }
+    }
+    return true
+  }
 }