Make Metadata Tests less Flaky (apache#10955)

Co-authored-by: Enrico Olivelli <eolivelli@apache.org>
yangl · Jun 23, 2021 · 71762a6 · 71762a6
1 parent 89e3986
commit 71762a6
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 11 deletions.
diff --git a/pulsar-metadata/pom.xml b/pulsar-metadata/pom.xml
@@ -31,7 +31,6 @@
 
   <artifactId>pulsar-metadata</artifactId>
   <name>Pulsar Metadata</name>
-
   <dependencies>
     <dependency>
       <groupId>org.apache.pulsar</groupId>
@@ -58,6 +57,12 @@
       <scope>test</scope>
     </dependency>
 
+    <dependency>
+      <groupId>org.awaitility</groupId>
+      <artifactId>awaitility</artifactId>
+      <scope>test</scope>
+    </dependency>
+
     <dependency>
       <groupId>org.apache.bookkeeper</groupId>
       <artifactId>bookkeeper-server</artifactId>

diff --git a/pulsar-metadata/src/test/java/org/apache/pulsar/metadata/ZKSessionTest.java b/pulsar-metadata/src/test/java/org/apache/pulsar/metadata/ZKSessionTest.java
@@ -26,7 +26,6 @@
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
-import java.util.function.Consumer;
 
 import lombok.Cleanup;
 
@@ -40,6 +39,7 @@
 import org.apache.pulsar.metadata.api.extended.SessionEvent;
 import org.apache.pulsar.metadata.coordination.impl.CoordinationServiceImpl;
 import org.apache.pulsar.metadata.impl.ZKMetadataStore;
+import org.awaitility.Awaitility;
 import org.testng.annotations.Test;
 
 public class ZKSessionTest extends BaseMetadataStoreTest {
@@ -61,10 +61,10 @@ public void testDisconnection() throws Exception {
         assertEquals(e, SessionEvent.ConnectionLost);
 
         zks.start();
-        e = sessionEvents.poll(10, TimeUnit.SECONDS);
+        e = sessionEvents.poll(20, TimeUnit.SECONDS);
         assertEquals(e, SessionEvent.Reconnected);
 
-        e = sessionEvents.poll(1, TimeUnit.SECONDS);
+        e = sessionEvents.poll(5, TimeUnit.SECONDS);
         assertNull(e);
     }
 
@@ -130,9 +130,9 @@ public void testReacquireLocksAfterSessionLost() throws Exception {
         e = sessionEvents.poll(10, TimeUnit.SECONDS);
         assertEquals(e, SessionEvent.SessionReestablished);
 
-        Thread.sleep(2_000);
-
-        assertFalse(lock.getLockExpiredFuture().isDone());
+        Awaitility.await().untilAsserted(() -> {
+            assertFalse(lock.getLockExpiredFuture().isDone());
+        });
 
         assertTrue(store.get(path).join().isPresent());
     }
@@ -171,7 +171,10 @@ public void testReacquireLeadershipAfterSessionLost() throws Exception {
         e = sessionEvents.poll(10, TimeUnit.SECONDS);
         assertEquals(e, SessionEvent.SessionLost);
 
-        assertEquals(le1.getState(), LeaderElectionState.Leading);
+        Awaitility.await().untilAsserted(() -> {
+            assertEquals(le1.getState(), LeaderElectionState.Leading);
+        });
+
         les = leaderElectionEvents.poll();
         assertNull(les);
 
@@ -180,9 +183,9 @@ public void testReacquireLeadershipAfterSessionLost() throws Exception {
         e = sessionEvents.poll(10, TimeUnit.SECONDS);
         assertEquals(e, SessionEvent.SessionReestablished);
 
-        Thread.sleep(2_000);
-
-        assertEquals(le1.getState(), LeaderElectionState.Leading);
+        Awaitility.await().untilAsserted(() -> {
+                    assertEquals(le1.getState(), LeaderElectionState.Leading);
+        });
         les = leaderElectionEvents.poll();
         assertNull(les);