diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/transaction/coordinator/TransactionMetaStoreAssignmentTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/transaction/coordinator/TransactionMetaStoreAssignmentTest.java
index cc0e185c582ba..ce173ba63d5d5 100644
--- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/transaction/coordinator/TransactionMetaStoreAssignmentTest.java
+++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/transaction/coordinator/TransactionMetaStoreAssignmentTest.java
@@ -20,6 +20,7 @@
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 import org.apache.pulsar.broker.PulsarService;
 import org.awaitility.Awaitility;
@@ -32,12 +33,13 @@ public class TransactionMetaStoreAssignmentTest extends TransactionMetaStoreTest
     @Test(groups = "broker")
     public void testTransactionMetaStoreAssignAndFailover() throws IOException {
 
-        int transactionMetaStoreCount = 0;
-        for (PulsarService pulsarService : pulsarServices) {
-            transactionMetaStoreCount += pulsarService.getTransactionMetadataStoreService().getStores().size();
-        }
-
-        Assert.assertEquals(transactionMetaStoreCount, 16);
+        Awaitility.await()
+                .untilAsserted(() -> {
+                    int transactionMetaStoreCount = Arrays.stream(pulsarServices)
+                            .mapToInt(pulsarService -> pulsarService.getTransactionMetadataStoreService().getStores().size())
+                            .sum();
+                    Assert.assertEquals(transactionMetaStoreCount, 16);
+                });
 
         PulsarService crashedMetaStore = null;
         for (int i = pulsarServices.length - 1; i >= 0; i--) {
@@ -62,16 +64,11 @@ public void testTransactionMetaStoreAssignAndFailover() throws IOException {
 
         Awaitility.await()
                 .untilAsserted(() -> {
-
-                    int transactionMetaStoreCount2 = 0;
-                    for (PulsarService pulsarService : pulsarServices) {
-                        transactionMetaStoreCount2 += pulsarService.getTransactionMetadataStoreService().getStores()
-                                .size();
-                    }
-
+                    int transactionMetaStoreCount2 = Arrays.stream(pulsarServices)
+                            .mapToInt(pulsarService -> pulsarService.getTransactionMetadataStoreService().getStores().size())
+                            .sum();
                     Assert.assertEquals(transactionMetaStoreCount2, 16);
                 });
-
         transactionCoordinatorClient.close();
     }
 }
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
index 24c4bd709007f..5f8c690752b4f 100644
--- a/pulsar-metadata/src/test/java/org/apache/pulsar/metadata/ZKSessionTest.java
+++ b/pulsar-metadata/src/test/java/org/apache/pulsar/metadata/ZKSessionTest.java
@@ -22,13 +22,10 @@
 import static org.testng.Assert.assertFalse;
 import static org.testng.Assert.assertNull;
 import static org.testng.Assert.assertTrue;
-
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
-
 import lombok.Cleanup;
-
 import org.apache.pulsar.metadata.api.MetadataStoreConfig;
 import org.apache.pulsar.metadata.api.coordination.CoordinationService;
 import org.apache.pulsar.metadata.api.coordination.LeaderElection;
@@ -139,6 +136,7 @@ public void testReacquireLocksAfterSessionLost() throws Exception {
 
     @Test
     public void testReacquireLeadershipAfterSessionLost() throws Exception {
+        //  ---  init
         @Cleanup
         MetadataStoreExtended store = MetadataStoreExtended.create(zks.getConnectionString(),
                 MetadataStoreConfig.builder()
@@ -147,7 +145,6 @@ public void testReacquireLeadershipAfterSessionLost() throws Exception {
 
         BlockingQueue<SessionEvent> sessionEvents = new LinkedBlockingQueue<>();
         store.registerSessionListener(sessionEvents::add);
-
         BlockingQueue<LeaderElectionState> leaderElectionEvents = new LinkedBlockingQueue<>();
         String path = newKey();
 
@@ -156,13 +153,13 @@ public void testReacquireLeadershipAfterSessionLost() throws Exception {
         @Cleanup
         LeaderElection<String> le1 = coordinationService.getLeaderElection(String.class, path,
                 leaderElectionEvents::add);
-
+        // --- test manual elect
         le1.elect("value-1").join();
         assertEquals(le1.getState(), LeaderElectionState.Leading);
 
         LeaderElectionState les = leaderElectionEvents.poll(5, TimeUnit.SECONDS);
         assertEquals(les, LeaderElectionState.Leading);
-
+        // --- expire session
         zks.expireSession(((ZKMetadataStore) store).getZkSessionId());
 
         SessionEvent e = sessionEvents.poll(5, TimeUnit.SECONDS);
@@ -170,25 +167,15 @@ public void testReacquireLeadershipAfterSessionLost() throws Exception {
 
         e = sessionEvents.poll(10, TimeUnit.SECONDS);
         assertEquals(e, SessionEvent.SessionLost);
-
-        Awaitility.await().untilAsserted(() -> {
-            assertEquals(le1.getState(), LeaderElectionState.Leading);
-        });
-
-        les = leaderElectionEvents.poll();
-        assertNull(les);
-
+        // --- test  le1 can be leader
+        Awaitility.await()
+                .untilAsserted(()-> assertEquals(le1.getState(),LeaderElectionState.Leading)); // reacquire leadership
         e = sessionEvents.poll(10, TimeUnit.SECONDS);
         assertEquals(e, SessionEvent.Reconnected);
         e = sessionEvents.poll(10, TimeUnit.SECONDS);
         assertEquals(e, SessionEvent.SessionReestablished);
-
-        Awaitility.await().untilAsserted(() -> {
-                    assertEquals(le1.getState(), LeaderElectionState.Leading);
-        });
-        les = leaderElectionEvents.poll();
-        assertNull(les);
-
+        Awaitility.await()
+                .untilAsserted(()-> assertEquals(le1.getState(),LeaderElectionState.Leading));
         assertTrue(store.get(path).join().isPresent());
     }
 }