test(unit-tests): improve reliability of message consumption in tests

Author: jeanouii

activemq-unit-tests/src/test/java/org/apache/activemq/store/StoreOrderTest.java

@@ -120,6 +120,7 @@ public void stopBroker() throws Exception {
         }
         if (broker != null) {
             broker.stop();
+            broker.waitUntilStopped();
         }
     }
 
@@ -257,6 +258,7 @@ protected BrokerService startBroker(boolean deleteMessagesOnStartup) throws Exce
         configureBroker(newBroker);
         newBroker.setDeleteAllMessagesOnStartup(deleteMessagesOnStartup);
         newBroker.start();
+        newBroker.waitUntilStarted();
         return newBroker;
     }
 

activemq-unit-tests/src/test/java/org/apache/activemq/store/jdbc/XACompletionTest.java

@@ -275,12 +275,7 @@ public void testStatsAndBrowseAfterAckPreparedRolledback() throws Exception {
 
         dumpMessages();
 
-        Wait.waitFor(new Wait.Condition() {
-            @Override
-            public boolean isSatisified() throws Exception {
-                return proxy.getInFlightCount() == 0l;
-            }
-        });
+        Wait.waitFor(() -> proxy.getInFlightCount() == 0L && proxy.cursorSize() == 0);
         assertEquals("prefetch", 0, proxy.getInFlightCount());
         assertEquals("size", 10, proxy.getQueueSize());
         assertEquals("cursor size", 0, proxy.cursorSize());

activemq-unit-tests/src/test/java/org/apache/activemq/usecases/QueueZeroPrefetchLazyDispatchPriorityTest.java

@@ -199,15 +199,23 @@ public void testPriorityMessagesWithJmsBrowser() throws Exception {
             assertNotNull(message);
             assertEquals(5, message.getJMSPriority());
 
-            // consume messages
-            final ArrayList<Message> consumeList = consumeMessages("TestQ");
+            // Wait for remaining messages to be fully available after consumeOneMessage closes its connection.
+            // With lazyDispatch=true + optimizedDispatch=true, messages may briefly be in-flight
+            // during connection teardown and not yet re-queued for dispatch to a new consumer.
+            final int remaining = numToSend - 1;
+            assertTrue("Remaining messages available for dispatch", Wait.waitFor(() -> {
+                final Queue q = (Queue) broker.getDestination(destination);
+                return q != null
+                    && q.getDestinationStatistics().getMessages().getCount() == remaining
+                    && q.getDestinationStatistics().getInflight().getCount() == 0;
+            }, 5000, 100));
+
+            // consume messages (use timeout-based overload for reliable dispatch on slow CI)
+            final ArrayList<Message> consumeList = consumeMessages("TestQ", remaining, TimeUnit.SECONDS.toMillis(30));
             LOG.info("Consumed list {}", consumeList.size());
 
-            // compare lists
-            // assertEquals("Iteration: " + i
-            // +", message 1 should be priority high", 5,
-            // consumeList.get(0).getJMSPriority());
-            for (int j = 1; j < (numToSend - 1); j++) {
+            assertEquals("Iteration: " + i + ", all remaining messages consumed", remaining, consumeList.size());
+            for (int j = 0; j < consumeList.size(); j++) {
                 assertEquals("Iteration: " + i + ", message " + j + " should be priority medium", 4, consumeList.get(j).getJMSPriority());
             }
         }