Another test timing issue

Author: kabir

extras/queue-manager-replicated/core/src/test/java/io/a2a/extras/queuemanager/replicated/core/ReplicatedQueueManagerTest.java

@@ -135,13 +135,10 @@ void testReplicatedEventDeliveredToCorrectQueue() throws InterruptedException {
         ReplicatedEventQueueItem replicatedEvent = new ReplicatedEventQueueItem(taskId, testEvent);
         queueManager.onReplicatedEvent(replicatedEvent);
 
-        Event dequeuedEvent;
-        try {
-            dequeuedEvent = queue.dequeueEventItem(100).getEvent();
-        } catch (EventQueueClosedException e) {
-            fail("Queue should not be closed");
-            return;
-        }
+        // Use retry logic to handle async MainEventBusProcessor distribution
+        EventQueueItem item = dequeueEventWithRetry(queue, 5000);
+        assertNotNull(item, "Event should be available in queue");
+        Event dequeuedEvent = item.getEvent();
         assertEquals(testEvent, dequeuedEvent);
     }
 
@@ -164,13 +161,11 @@ void testReplicatedEventCreatesQueueIfNeeded() throws InterruptedException {
         // Verify the event was enqueued by dequeuing it
         // Need to tap() the MainQueue to get a ChildQueue for consumption
         EventQueue childQueue = queue.tap();
-        Event dequeuedEvent;
-        try {
-            dequeuedEvent = childQueue.dequeueEventItem(100).getEvent();
-        } catch (EventQueueClosedException e) {
-            fail("Queue should not be closed");
-            return;
-        }
+
+        // Use retry logic to handle async MainEventBusProcessor distribution
+        EventQueueItem item = dequeueEventWithRetry(childQueue, 5000);
+        assertNotNull(item, "Event should be available in queue");
+        Event dequeuedEvent = item.getEvent();
         assertEquals(testEvent, dequeuedEvent, "The replicated event should be enqueued in the newly created queue");
     }
 
@@ -355,27 +350,24 @@ void testReplicatedEventProcessedWhenTaskActive() throws InterruptedException {
         // Verify no queue exists initially
         assertNull(queueManager.get(taskId));
 
+        // Create a ChildQueue BEFORE processing the replicated event
+        // This ensures the ChildQueue exists when MainEventBusProcessor distributes the event
+        EventQueue childQueue = queueManager.createOrTap(taskId);
+        assertNotNull(childQueue, "ChildQueue should be created");
+
+        // Verify MainQueue was created
+        EventQueue mainQueue = queueManager.get(taskId);
+        assertNotNull(mainQueue, "MainQueue should exist after createOrTap");
+
         // Process a replicated event for an active task
         ReplicatedEventQueueItem replicatedEvent = new ReplicatedEventQueueItem(taskId, testEvent);
         queueManager.onReplicatedEvent(replicatedEvent);
 
-        // Queue should be created and event should be enqueued
-        EventQueue queue = queueManager.get(taskId);
-        assertNotNull(queue, "Queue should be created for active task");
-
-        // Verify the event was enqueued
-        // Need to tap() the MainQueue to get a ChildQueue for consumption
-        EventQueue childQueue = queue.tap();
-        Event dequeuedEvent;
-        try {
-            // Wait up to 5 seconds for event (async processing needs time)
-            EventQueueItem item = childQueue.dequeueEventItem(5000);
-            assertNotNull(item, "Event should be available in queue");
-            dequeuedEvent = item.getEvent();
-        } catch (EventQueueClosedException e) {
-            fail("Queue should not be closed");
-            return;
-        }
+        // Verify the event was enqueued and distributed to our ChildQueue
+        // Use retry logic to handle async MainEventBusProcessor distribution
+        EventQueueItem item = dequeueEventWithRetry(childQueue, 5000);
+        assertNotNull(item, "Event should be available in queue");
+        Event dequeuedEvent = item.getEvent();
         assertEquals(testEvent, dequeuedEvent, "Event should be enqueued for active task");
     }
 
@@ -498,26 +490,14 @@ void testReplicatedQueueClosedEventTerminatesConsumer() throws InterruptedExcept
         ReplicatedEventQueueItem replicatedClosedEvent = new ReplicatedEventQueueItem(taskId, closedEvent);
         queueManager.onReplicatedEvent(replicatedClosedEvent);
 
-        // Dequeue the normal event first
-        EventQueueItem item1;
-        try {
-            item1 = queue.dequeueEventItem(100);
-        } catch (EventQueueClosedException e) {
-            fail("Should not throw on first dequeue");
-            return;
-        }
-        assertNotNull(item1);
+        // Dequeue the normal event first (use retry for async processing)
+        EventQueueItem item1 = dequeueEventWithRetry(queue, 5000);
+        assertNotNull(item1, "First event should be available");
         assertEquals(testEvent, item1.getEvent());
 
-        // Next dequeue should get the QueueClosedEvent
-        EventQueueItem item2;
-        try {
-            item2 = queue.dequeueEventItem(100);
-        } catch (EventQueueClosedException e) {
-            fail("Should not throw on second dequeue, should return the event");
-            return;
-        }
-        assertNotNull(item2);
+        // Next dequeue should get the QueueClosedEvent (use retry for async processing)
+        EventQueueItem item2 = dequeueEventWithRetry(queue, 5000);
+        assertNotNull(item2, "QueueClosedEvent should be available");
         assertTrue(item2.getEvent() instanceof QueueClosedEvent,
                 "Second event should be QueueClosedEvent");
     }
@@ -576,4 +556,30 @@ public void setActive(boolean active) {
             this.active = active;
         }
     }
+
+    /**
+     * Helper method to dequeue an event with retry logic to handle async MainEventBusProcessor distribution.
+     * This handles the race condition where MainEventBusProcessor may not have distributed the event yet.
+     *
+     * @param queue the queue to dequeue from
+     * @param maxWaitMs maximum time to wait in milliseconds (default 5000)
+     * @return the dequeued EventQueueItem, or null if timeout occurs
+     */
+    private EventQueueItem dequeueEventWithRetry(EventQueue queue, long maxWaitMs) throws InterruptedException {
+        EventQueueItem item = null;
+        long pollStart = System.currentTimeMillis();
+        while (item == null && (System.currentTimeMillis() - pollStart) < maxWaitMs) {
+            try {
+                item = queue.dequeueEventItem(100); // Short timeout per attempt
+                if (item == null) {
+                    // No event yet, wait a bit before retry
+                    Thread.sleep(50);
+                }
+            } catch (EventQueueClosedException e) {
+                // Queue closed, return null
+                return null;
+            }
+        }
+        return item;
+    }
 }