[fix][flaky-test]ProxyConnectionThrottlingTest.testInboundConnection (#…

…17883)
apache · Sep 29, 2022 · 3de690d · 3de690d
1 parent 0f53c79
commit 3de690d
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 4 deletions.
diff --git a/pulsar-proxy/pom.xml b/pulsar-proxy/pom.xml
@@ -159,6 +159,12 @@
       <scope>test</scope>
     </dependency>
 
+    <dependency>
+      <groupId>org.awaitility</groupId>
+      <artifactId>awaitility</artifactId>
+      <scope>test</scope>
+    </dependency>
+
     <dependency>
       <groupId>com.beust</groupId>
       <artifactId>jcommander</artifactId>

diff --git a/pulsar-proxy/src/test/java/org/apache/pulsar/proxy/server/ProxyConnectionThrottlingTest.java b/pulsar-proxy/src/test/java/org/apache/pulsar/proxy/server/ProxyConnectionThrottlingTest.java
@@ -32,6 +32,7 @@
 import org.apache.pulsar.client.api.Schema;
 import org.apache.pulsar.common.configuration.PulsarConfigurationLoader;
 import org.apache.pulsar.metadata.impl.ZKMetadataStore;
+import org.awaitility.Awaitility;
 import org.mockito.Mockito;
 import org.testng.Assert;
 import org.testng.annotations.AfterClass;
@@ -108,7 +109,9 @@ public void testInboundConnection() throws Exception {
         } catch (Exception ex) {
             // OK
         }
-        Assert.assertEquals(ConnectionController.DefaultConnectionController.getTotalConnectionNum(), 4);
+        Awaitility.await().untilAsserted(() ->{
+            Assert.assertEquals(ConnectionController.DefaultConnectionController.getTotalConnectionNum(), 4);
+        });
         Assert.assertEquals(ConnectionController.DefaultConnectionController.getConnections().size(), 1);
         Set<String> keys = ConnectionController.DefaultConnectionController.getConnections().keySet();
         for (String key : keys) {
@@ -119,7 +122,9 @@ public void testInboundConnection() throws Exception {
 
         client1.close();
 
-        Assert.assertEquals(ConnectionController.DefaultConnectionController.getTotalConnectionNum(), 2);
+        Awaitility.await().untilAsserted(() ->{
+            Assert.assertEquals(ConnectionController.DefaultConnectionController.getTotalConnectionNum(), 2);
+        });
         Assert.assertEquals(ConnectionController.DefaultConnectionController.getConnections().size(), 1);
         keys = ConnectionController.DefaultConnectionController.getConnections().keySet();
         for (String key : keys) {
@@ -129,8 +134,9 @@ public void testInboundConnection() throws Exception {
         Assert.assertEquals(ProxyService.ACTIVE_CONNECTIONS.get(), 2.0d);
 
         client2.close();
-
-        Assert.assertEquals(ConnectionController.DefaultConnectionController.getTotalConnectionNum(), 0);
+        Awaitility.await().untilAsserted(() ->{
+            Assert.assertEquals(ConnectionController.DefaultConnectionController.getTotalConnectionNum(), 0);
+        });
         Assert.assertEquals(ConnectionController.DefaultConnectionController.getConnections().size(), 0);
         Assert.assertEquals(ProxyService.ACTIVE_CONNECTIONS.get(), 0.0d);
     }