Fix dead lettering flaky test

It was using Thread.sleep() statements, which are not reliable,
using polling for assertion instead.

(cherry picked from commit ea8aa89)

Author: acogoluegnes

src/test/java/com/rabbitmq/client/test/TestUtils.java

@@ -56,9 +56,20 @@ public static ConnectionFactory connectionFactory() {
         return connectionFactory;
     }
 
-    public static void waitAtMost(Duration timeout, BooleanSupplier condition) {
-        if (condition.getAsBoolean()) {
-            return;
+    @FunctionalInterface
+    public interface CallableBooleanSupplier {
+
+        boolean getAsBoolean() throws Exception;
+
+    }
+
+    public static void waitAtMost(Duration timeout, CallableBooleanSupplier condition) {
+        try {
+            if (condition.getAsBoolean()) {
+                return;
+            }
+        } catch (Exception e) {
+            throw new RuntimeException(e);
         }
         int waitTime = 100;
         int waitedTime = 0;
@@ -70,8 +81,12 @@ public static void waitAtMost(Duration timeout, BooleanSupplier condition) {
                 Thread.currentThread().interrupt();
                 throw new RuntimeException(e);
             }
-            if (condition.getAsBoolean()) {
-                return;
+            try {
+                if (condition.getAsBoolean()) {
+                    return;
+                }
+            } catch (Exception e) {
+                throw new RuntimeException(e);
             }
             waitedTime += waitTime;
         }

src/test/java/com/rabbitmq/client/test/functional/DeadLetterExchange.java

@@ -19,12 +19,15 @@
 import com.rabbitmq.client.AMQP.BasicProperties;
 import com.rabbitmq.client.test.BrokerTestCase;
 import com.rabbitmq.client.test.TestUtils;
+import java.util.concurrent.atomic.AtomicReference;
 import org.junit.Test;
 
 import java.io.IOException;
 import java.util.*;
 import java.util.concurrent.*;
 
+import static com.rabbitmq.client.test.TestUtils.waitAtMost;
+import static java.time.Duration.ofSeconds;
 import static org.junit.Assert.*;
 
 public class DeadLetterExchange extends BrokerTestCase {
@@ -403,9 +406,15 @@ public void handleDelivery(String consumerTag, Envelope envelope,
         channel.queueBind(DLQ, DLX, "test");
         publishN(1);
 
-        sleep(200);
-
-        GetResponse getResponse = channel.basicGet(DLQ, true);
+        AtomicReference<GetResponse> responseRefeference = new AtomicReference<>();
+        waitAtMost(
+            ofSeconds(1),
+            () -> {
+                GetResponse response = channel.basicGet(DLQ, true);
+                responseRefeference.set(response);
+                return responseRefeference.get() != null;
+        });
+        GetResponse getResponse = responseRefeference.get();
         assertNotNull("Message not dead-lettered",
             getResponse);
         assertEquals("test message", new String(getResponse.getBody()));
@@ -432,9 +441,15 @@ public void handleDelivery(String consumerTag, Envelope envelope,
                .headers(headers)
                .build(), "test message".getBytes());
 
-        sleep(100);
-
-        getResponse = channel.basicGet(DLQ, true);
+        responseRefeference.set(null);
+        waitAtMost(
+            ofSeconds(1),
+            () -> {
+                GetResponse response = channel.basicGet(DLQ, true);
+                responseRefeference.set(response);
+                return responseRefeference.get() != null;
+            });
+        getResponse = responseRefeference.get();
         assertNotNull("Message not dead-lettered", getResponse);
         assertEquals("test message", new String(getResponse.getBody()));
         headers = getResponse.getProps().getHeaders();
@@ -453,9 +468,17 @@ public void handleDelivery(String consumerTag, Envelope envelope,
             new AMQP.BasicProperties.Builder()
                .headers(headers)
                .build(), "test message".getBytes());
-        sleep(100);
 
-        getResponse = channel.basicGet(DLQ, true);
+        responseRefeference.set(null);
+        waitAtMost(
+            ofSeconds(1),
+            () -> {
+                GetResponse response = channel.basicGet(DLQ, true);
+                responseRefeference.set(response);
+                return responseRefeference.get() != null;
+            });
+        getResponse = responseRefeference.get();
+
         assertNotNull("Message not dead-lettered", getResponse);
         assertEquals("test message", new String(getResponse.getBody()));
         headers = getResponse.getProps().getHeaders();