From 22109b2719d35f15fba7c58095104e12594fe717 Mon Sep 17 00:00:00 2001
From: Lauri Tulmin <ltulmin@splunk.com>
Date: Sat, 6 Feb 2021 02:06:09 +0200
Subject: [PATCH] Fix flaky spring jms test (#2202)

---
 .../jms2Test/groovy/SpringTemplateJms2Test.groovy | 15 +++++++++++----
 .../src/test/groovy/SpringTemplateJms1Test.groovy | 11 +++++++++++
 2 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/instrumentation/jms-1.1/javaagent/src/jms2Test/groovy/SpringTemplateJms2Test.groovy b/instrumentation/jms-1.1/javaagent/src/jms2Test/groovy/SpringTemplateJms2Test.groovy
index 439c2affe0..23110488d5 100644
--- a/instrumentation/jms-1.1/javaagent/src/jms2Test/groovy/SpringTemplateJms2Test.groovy
+++ b/instrumentation/jms-1.1/javaagent/src/jms2Test/groovy/SpringTemplateJms2Test.groovy
@@ -8,7 +8,6 @@ import static Jms2Test.producerSpan
 
 import com.google.common.io.Files
 import io.opentelemetry.instrumentation.test.AgentInstrumentationSpecification
-import java.util.concurrent.CountDownLatch
 import java.util.concurrent.TimeUnit
 import java.util.concurrent.atomic.AtomicReference
 import javax.jms.Session
@@ -102,9 +101,7 @@ class SpringTemplateJms2Test extends AgentInstrumentationSpecification {
   def "send and receive message generates spans"() {
     setup:
     AtomicReference<String> msgId = new AtomicReference<>()
-    CountDownLatch countDownLatch = new CountDownLatch(1)
     Thread.start {
-      countDownLatch.countDown()
       TextMessage msg = template.receive(destination)
       assert msg.text == messageText
       msgId.set(msg.getJMSMessageID())
@@ -115,7 +112,6 @@ class SpringTemplateJms2Test extends AgentInstrumentationSpecification {
       }
     }
     // wait for thread to start, we expect the first span to be from receive
-    countDownLatch.await()
     TextMessage receivedMessage = template.sendAndReceive(destination) {
       session -> template.getMessageConverter().toMessage(messageText, session)
     }
@@ -123,6 +119,17 @@ class SpringTemplateJms2Test extends AgentInstrumentationSpecification {
     expect:
     receivedMessage.text == "responded!"
     assertTraces(4) {
+      sortTraces {
+        def expectedOrder = ["$destinationName receive",
+                             "$destinationName send",
+                             "(temporary) receive",
+                             "(temporary) send"]
+        // ensure that traces appear in expected order
+        traces.sort {a,b ->
+          expectedOrder.indexOf(a[0].name) - expectedOrder.indexOf(b[0].name)
+        }
+      }
+
       trace(0, 1) {
         consumerSpan(it, 0, destinationType, destinationName, msgId.get(), null, "receive")
       }
diff --git a/instrumentation/jms-1.1/javaagent/src/test/groovy/SpringTemplateJms1Test.groovy b/instrumentation/jms-1.1/javaagent/src/test/groovy/SpringTemplateJms1Test.groovy
index 730e941920..8c990d7e44 100644
--- a/instrumentation/jms-1.1/javaagent/src/test/groovy/SpringTemplateJms1Test.groovy
+++ b/instrumentation/jms-1.1/javaagent/src/test/groovy/SpringTemplateJms1Test.groovy
@@ -97,6 +97,17 @@ class SpringTemplateJms1Test extends AgentInstrumentationSpecification {
     expect:
     receivedMessage.text == "responded!"
     assertTraces(4) {
+      sortTraces {
+        def expectedOrder = ["$destinationName receive",
+                             "$destinationName send",
+                             "(temporary) receive",
+                             "(temporary) send"]
+        // ensure that traces appear in expected order
+        traces.sort {a,b ->
+          expectedOrder.indexOf(a[0].name) - expectedOrder.indexOf(b[0].name)
+        }
+      }
+
       trace(0, 1) {
         consumerSpan(it, 0, destinationType, destinationName, msgId.get(), null, "receive")
       }