From a862b1511c7df151a7c09cec59761501a4c274c2 Mon Sep 17 00:00:00 2001
From: eyalkoren <eyal.koren@elastic.co>
Date: Wed, 12 Jun 2019 07:48:44 +0300
Subject: [PATCH 1/2] Increase timeouts to 3 seconds

---
 .../co/elastic/apm/agent/jms/JmsInstrumentationIT.java | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/apm-agent-plugins/apm-jms-plugin/src/test/java/co/elastic/apm/agent/jms/JmsInstrumentationIT.java b/apm-agent-plugins/apm-jms-plugin/src/test/java/co/elastic/apm/agent/jms/JmsInstrumentationIT.java
index 6e5aa0e633..4865ffac58 100644
--- a/apm-agent-plugins/apm-jms-plugin/src/test/java/co/elastic/apm/agent/jms/JmsInstrumentationIT.java
+++ b/apm-agent-plugins/apm-jms-plugin/src/test/java/co/elastic/apm/agent/jms/JmsInstrumentationIT.java
@@ -141,7 +141,7 @@ private void testQueueSendReceiveOnTracedThread(Callable<Message> receiveMethod,
         String message = UUID.randomUUID().toString();
         Message outgoingMessage = brokerFacade.createTextMessage(message);
         brokerFacade.send(queue, outgoingMessage);
-        Message incomingMessage = receiveOnTracedThread(receiveMethod).get(1, TimeUnit.SECONDS);
+        Message incomingMessage = receiveOnTracedThread(receiveMethod).get(3, TimeUnit.SECONDS);
         verifyMessage(message, incomingMessage);
 
         List<Span> spans = reporter.getSpans();
@@ -168,7 +168,7 @@ private void testQueueSendReceiveOnNonTracedThread(Callable<Message> receiveMeth
         String message = UUID.randomUUID().toString();
         Message outgoingMessage = brokerFacade.createTextMessage(message);
         brokerFacade.send(queue, outgoingMessage);
-        Message incomingMessage = receiveOnNonTracedThread(receiveMethod).get(1, TimeUnit.SECONDS);
+        Message incomingMessage = receiveOnNonTracedThread(receiveMethod).get(3, TimeUnit.SECONDS);
         verifyMessage(message, incomingMessage);
         verifySendReceiveOnNonTracedThread(queue.getQueueName(), 1);
     }
@@ -228,7 +228,7 @@ private void testQueueSendListen(Function<Destination, CompletableFuture<Message
         String message = UUID.randomUUID().toString();
         Message outgoingMessage = brokerFacade.createTextMessage(message);
         brokerFacade.send(queue, outgoingMessage);
-        Message incomingMessage = incomingMessageFuture.get(1, TimeUnit.SECONDS);
+        Message incomingMessage = incomingMessageFuture.get(3, TimeUnit.SECONDS);
         verifyMessage(message, incomingMessage);
         Thread.sleep(500);
         verifySendReceiveOnNonTracedThread(queue.getQueueName(), 1);
@@ -245,10 +245,10 @@ public void testTopicWithTwoSubscribers() throws Exception {
         Message outgoingMessage = brokerFacade.createTextMessage(message);
         brokerFacade.send(topic, outgoingMessage);
 
-        Message incomingMessage1 = incomingMessageFuture1.get(1, TimeUnit.SECONDS);
+        Message incomingMessage1 = incomingMessageFuture1.get(3, TimeUnit.SECONDS);
         verifyMessage(message, incomingMessage1);
 
-        Message incomingMessage2 = incomingMessageFuture2.get(1, TimeUnit.SECONDS);
+        Message incomingMessage2 = incomingMessageFuture2.get(3, TimeUnit.SECONDS);
         verifyMessage(message, incomingMessage2);
 
         Thread.sleep(500);

From 7399079d9cceeb060a25ab954f7f8282991fe53c Mon Sep 17 00:00:00 2001
From: eyalkoren <eyal.koren@elastic.co>
Date: Wed, 12 Jun 2019 13:23:47 +0300
Subject: [PATCH 2/2] Adding a loop on receiveNoWait API test

---
 .../apm/agent/jms/JmsInstrumentationIT.java   | 21 +++++++++++++++++--
 1 file changed, 19 insertions(+), 2 deletions(-)

diff --git a/apm-agent-plugins/apm-jms-plugin/src/test/java/co/elastic/apm/agent/jms/JmsInstrumentationIT.java b/apm-agent-plugins/apm-jms-plugin/src/test/java/co/elastic/apm/agent/jms/JmsInstrumentationIT.java
index 4865ffac58..53af2300b5 100644
--- a/apm-agent-plugins/apm-jms-plugin/src/test/java/co/elastic/apm/agent/jms/JmsInstrumentationIT.java
+++ b/apm-agent-plugins/apm-jms-plugin/src/test/java/co/elastic/apm/agent/jms/JmsInstrumentationIT.java
@@ -56,6 +56,7 @@
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import java.util.function.Function;
 
 import static org.assertj.core.api.Assertions.assertThat;
@@ -122,7 +123,7 @@ public void testQueueSendReceiveOnTracedThread() throws Exception {
     @Test
     public void testQueueSendReceiveNoWaitOnTracedThread() throws Exception {
         final Queue queue = brokerFacade.createQueue(TEST_Q_NAME);
-        testQueueSendReceiveOnTracedThread(() -> brokerFacade.receiveNoWait(queue), queue);
+        testQueueSendReceiveOnTracedThread(() -> loopReceive(() -> brokerFacade.receiveNoWait(queue), 3000), queue);
     }
 
     @Test
@@ -134,7 +135,23 @@ public void testQueueSendReceiveOnNonTracedThread() throws Exception {
     @Test
     public void testQueueSendReceiveNoWaitOnNonTracedThread() throws Exception {
         final Queue queue = brokerFacade.createQueue(TEST_Q_NAME);
-        testQueueSendReceiveOnNonTracedThread(() -> brokerFacade.receiveNoWait(queue), queue);
+        testQueueSendReceiveOnNonTracedThread(() -> loopReceive(() -> brokerFacade.receiveNoWait(queue), 3000), queue);
+    }
+
+    // A utility method for testing the receiveNoWait API consistently
+    private Message loopReceive(Callable<Message> receiveMethod, @SuppressWarnings("SameParameterValue") long timeout) throws Exception {
+        long start = System.currentTimeMillis();
+        long curr = start;
+        while (curr - start < timeout) {
+            Message ret = receiveMethod.call();
+            if (ret != null) {
+                return ret;
+            } else {
+                Thread.sleep(100);
+                curr = System.currentTimeMillis();
+            }
+        }
+        throw new TimeoutException();
     }
 
     private void testQueueSendReceiveOnTracedThread(Callable<Message> receiveMethod, Queue queue) throws Exception {