Use callback rather than polling where possible in tests

Author: kabir

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

@@ -83,13 +83,45 @@ private TaskStatusUpdateEvent createEventForTask(String taskId) {
     @AfterEach
     void tearDown() {
         if (mainEventBusProcessor != null) {
+            mainEventBusProcessor.setCallback(null);  // Clear any test callbacks
             EventQueueUtil.stop(mainEventBusProcessor);
         }
         mainEventBusProcessor = null;
         mainEventBus = null;
         queueManager = null;
     }
 
+    /**
+     * Helper to wait for MainEventBusProcessor to process an event.
+     * Replaces polling patterns with deterministic callback-based waiting.
+     *
+     * @param action the action that triggers event processing
+     * @throws InterruptedException if waiting is interrupted
+     * @throws AssertionError if processing doesn't complete within timeout
+     */
+    private void waitForEventProcessing(Runnable action) throws InterruptedException {
+        CountDownLatch processingLatch = new CountDownLatch(1);
+        mainEventBusProcessor.setCallback(new io.a2a.server.events.MainEventBusProcessorCallback() {
+            @Override
+            public void onEventProcessed(String taskId, io.a2a.spec.Event event) {
+                processingLatch.countDown();
+            }
+
+            @Override
+            public void onTaskFinalized(String taskId) {
+                // Not needed for basic event processing wait
+            }
+        });
+
+        try {
+            action.run();
+            assertTrue(processingLatch.await(5, TimeUnit.SECONDS),
+                    "MainEventBusProcessor should have processed the event within timeout");
+        } finally {
+            mainEventBusProcessor.setCallback(null);
+        }
+    }
+
     @Test
     void testReplicationStrategyTriggeredOnNormalEnqueue() throws InterruptedException {
         CountingReplicationStrategy strategy = new CountingReplicationStrategy();
@@ -147,10 +179,9 @@ void testReplicatedEventDeliveredToCorrectQueue() throws InterruptedException {
         EventQueue queue = queueManager.createOrTap(taskId);
 
         ReplicatedEventQueueItem replicatedEvent = new ReplicatedEventQueueItem(taskId, eventForTask);
-        queueManager.onReplicatedEvent(replicatedEvent);
 
-        // Use retry logic to handle async MainEventBusProcessor distribution
-        EventQueueItem item = dequeueEventWithRetry(queue, 5000);
+        // Use callback to wait for event processing
+        EventQueueItem item = dequeueEventWithRetry(queue, () -> queueManager.onReplicatedEvent(replicatedEvent));
         assertNotNull(item, "Event should be available in queue");
         Event dequeuedEvent = item.getEvent();
         assertEquals(eventForTask, dequeuedEvent);
@@ -175,12 +206,11 @@ void testReplicatedEventCreatesQueueIfNeeded() throws InterruptedException {
 
         ReplicatedEventQueueItem replicatedEvent = new ReplicatedEventQueueItem(taskId, eventForTask);
 
-        // Process the replicated event
-        assertDoesNotThrow(() -> queueManager.onReplicatedEvent(replicatedEvent));
-
-        // Verify the event was enqueued and distributed to our ChildQueue
-        // Use retry logic to handle async MainEventBusProcessor distribution
-        EventQueueItem item = dequeueEventWithRetry(childQueue, 5000);
+        // Process the replicated event and wait for distribution
+        // Use callback to wait for event processing
+        EventQueueItem item = dequeueEventWithRetry(childQueue, () -> {
+            assertDoesNotThrow(() -> queueManager.onReplicatedEvent(replicatedEvent));
+        });
         assertNotNull(item, "Event should be available in queue");
         Event dequeuedEvent = item.getEvent();
         assertEquals(eventForTask, dequeuedEvent, "The replicated event should be enqueued in the newly created queue");
@@ -379,11 +409,10 @@ void testReplicatedEventProcessedWhenTaskActive() throws InterruptedException {
 
         // Process a replicated event for an active task
         ReplicatedEventQueueItem replicatedEvent = new ReplicatedEventQueueItem(taskId, eventForTask);
-        queueManager.onReplicatedEvent(replicatedEvent);
 
         // Verify the event was enqueued and distributed to our ChildQueue
-        // Use retry logic to handle async MainEventBusProcessor distribution
-        EventQueueItem item = dequeueEventWithRetry(childQueue, 5000);
+        // Use callback to wait for event processing
+        EventQueueItem item = dequeueEventWithRetry(childQueue, () -> queueManager.onReplicatedEvent(replicatedEvent));
         assertNotNull(item, "Event should be available in queue");
         Event dequeuedEvent = item.getEvent();
         assertEquals(eventForTask, dequeuedEvent, "Event should be enqueued for active task");
@@ -501,21 +530,17 @@ void testReplicatedQueueClosedEventTerminatesConsumer() throws InterruptedExcept
         TaskStatusUpdateEvent eventForTask = createEventForTask(taskId);  // Use matching taskId
         EventQueue queue = queueManager.createOrTap(taskId);
 
-        // Enqueue a normal event
-        queue.enqueueEvent(eventForTask);
-
         // Simulate receiving QueueClosedEvent from remote node
         QueueClosedEvent closedEvent = new QueueClosedEvent(taskId);
         ReplicatedEventQueueItem replicatedClosedEvent = new ReplicatedEventQueueItem(taskId, closedEvent);
-        queueManager.onReplicatedEvent(replicatedClosedEvent);
 
-        // Dequeue the normal event first (use retry for async processing)
-        EventQueueItem item1 = dequeueEventWithRetry(queue, 5000);
+        // Dequeue the normal event first (use callback to wait for async processing)
+        EventQueueItem item1 = dequeueEventWithRetry(queue, () -> queue.enqueueEvent(eventForTask));
         assertNotNull(item1, "First event should be available");
         assertEquals(eventForTask, item1.getEvent());
 
-        // Next dequeue should get the QueueClosedEvent (use retry for async processing)
-        EventQueueItem item2 = dequeueEventWithRetry(queue, 5000);
+        // Next dequeue should get the QueueClosedEvent (use callback to wait for async processing)
+        EventQueueItem item2 = dequeueEventWithRetry(queue, () -> queueManager.onReplicatedEvent(replicatedClosedEvent));
         assertNotNull(item2, "QueueClosedEvent should be available");
         assertTrue(item2.getEvent() instanceof QueueClosedEvent,
                 "Second event should be QueueClosedEvent");
@@ -577,28 +602,23 @@ public void setActive(boolean 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.
+     * Helper method to dequeue an event after waiting for MainEventBusProcessor distribution.
+     * Uses callback-based waiting instead of polling for deterministic synchronization.
      *
      * @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
+     * @param enqueueAction the action that enqueues the event (triggers event processing)
+     * @return the dequeued EventQueueItem, or null if queue is closed
      */
-    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;
-            }
+    private EventQueueItem dequeueEventWithRetry(EventQueue queue, Runnable enqueueAction) throws InterruptedException {
+        // Wait for event to be processed and distributed
+        waitForEventProcessing(enqueueAction);
+
+        // Event is now available, dequeue directly
+        try {
+            return queue.dequeueEventItem(100);
+        } catch (EventQueueClosedException e) {
+            // Queue closed, return null
+            return null;
         }
-        return item;
     }
 }

server-common/src/test/java/io/a2a/server/events/EventConsumerTest.java

@@ -80,10 +80,42 @@ public void init() {
     @AfterEach
     public void cleanup() {
         if (mainEventBusProcessor != null) {
+            mainEventBusProcessor.setCallback(null);  // Clear any test callbacks
             EventQueueUtil.stop(mainEventBusProcessor);
         }
     }
 
+    /**
+     * Helper to wait for MainEventBusProcessor to process an event.
+     * Replaces polling patterns with deterministic callback-based waiting.
+     *
+     * @param action the action that triggers event processing
+     * @throws InterruptedException if waiting is interrupted
+     * @throws AssertionError if processing doesn't complete within timeout
+     */
+    private void waitForEventProcessing(Runnable action) throws InterruptedException {
+        CountDownLatch processingLatch = new CountDownLatch(1);
+        mainEventBusProcessor.setCallback(new MainEventBusProcessorCallback() {
+            @Override
+            public void onEventProcessed(String taskId, Event event) {
+                processingLatch.countDown();
+            }
+
+            @Override
+            public void onTaskFinalized(String taskId) {
+                // Not needed for basic event processing wait
+            }
+        });
+
+        try {
+            action.run();
+            assertTrue(processingLatch.await(5, TimeUnit.SECONDS),
+                    "MainEventBusProcessor should have processed the event within timeout");
+        } finally {
+            mainEventBusProcessor.setCallback(null);
+        }
+    }
+
     @Test
     public void testConsumeOneTaskEvent() throws Exception {
         Task event = Utils.unmarshalFrom(MINIMAL_TASK, Task.class);
@@ -421,15 +453,9 @@ public void testConsumeAllHandlesQueueClosedException() throws Exception {
 
         // Add a message event (which will complete the stream)
         Event message = Utils.unmarshalFrom(MESSAGE_PAYLOAD, Message.class);
-        queue.enqueueEvent(message);
 
-        // Poll for event to arrive in ChildQueue (async MainEventBusProcessor distribution)
-        long startTime = System.currentTimeMillis();
-        long timeout = 2000;
-        while (queue.size() == 0 && (System.currentTimeMillis() - startTime) < timeout) {
-            Thread.sleep(50);
-        }
-        assertTrue(queue.size() > 0, "Event should arrive in ChildQueue within timeout");
+        // Use callback to wait for event processing
+        waitForEventProcessing(() -> queue.enqueueEvent(message));
 
         // Close the queue before consuming
         queue.close();
@@ -525,20 +551,12 @@ public void onComplete() {
     }
 
     private void enqueueAndConsumeOneEvent(Event event) throws Exception {
-        eventQueue.enqueueEvent(event);
-        // Poll for event with 2-second timeout
-        long startTime = System.currentTimeMillis();
-        long timeout = 2000;
-        Event result = null;
-        while (result == null && (System.currentTimeMillis() - startTime) < timeout) {
-            try {
-                result = eventConsumer.consumeOne();
-            } catch (A2AServerException e) {
-                // Event not available yet, wait a bit and try again
-                Thread.sleep(50);
-            }
-        }
-        assertNotNull(result, "Event should arrive within timeout");
+        // Use callback to wait for event processing
+        waitForEventProcessing(() -> eventQueue.enqueueEvent(event));
+
+        // Event is now available, consume it directly
+        Event result = eventConsumer.consumeOne();
+        assertNotNull(result, "Event should be available");
         assertSame(event, result);
     }
 

server-common/src/test/java/io/a2a/server/events/EventQueueTest.java

@@ -10,6 +10,8 @@
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 
 import io.a2a.server.tasks.InMemoryTaskStore;
 import io.a2a.server.tasks.PushNotificationSender;
@@ -74,6 +76,7 @@ public void init() {
     @AfterEach
     public void cleanup() {
         if (mainEventBusProcessor != null) {
+            mainEventBusProcessor.setCallback(null);  // Clear any test callbacks
             EventQueueUtil.stop(mainEventBusProcessor);
         }
     }
@@ -87,6 +90,37 @@ private EventQueue createQueueWithEventBus(String taskId) {
                 .build();
     }
 
+    /**
+     * Helper to wait for MainEventBusProcessor to process an event.
+     * Replaces polling patterns with deterministic callback-based waiting.
+     *
+     * @param action the action that triggers event processing
+     * @throws InterruptedException if waiting is interrupted
+     * @throws AssertionError if processing doesn't complete within timeout
+     */
+    private void waitForEventProcessing(Runnable action) throws InterruptedException {
+        CountDownLatch processingLatch = new CountDownLatch(1);
+        mainEventBusProcessor.setCallback(new io.a2a.server.events.MainEventBusProcessorCallback() {
+            @Override
+            public void onEventProcessed(String taskId, io.a2a.spec.Event event) {
+                processingLatch.countDown();
+            }
+
+            @Override
+            public void onTaskFinalized(String taskId) {
+                // Not needed for basic event processing wait
+            }
+        });
+
+        try {
+            action.run();
+            assertTrue(processingLatch.await(5, TimeUnit.SECONDS),
+                    "MainEventBusProcessor should have processed the event within timeout");
+        } finally {
+            mainEventBusProcessor.setCallback(null);
+        }
+    }
+
     @Test
     public void testConstructorDefaultQueueSize() {
         EventQueue queue = EventQueueUtil.getEventQueueBuilder(mainEventBus).build();
@@ -269,17 +303,10 @@ public void testDequeueEventWhenClosedButHasEvents() throws Exception {
         EventQueue childQueue = mainQueue.tap();
         Event event = Utils.unmarshalFrom(MINIMAL_TASK, Task.class);
 
-        // Enqueue to mainQueue
-        mainQueue.enqueueEvent(event);
-
-        // Poll for event to arrive in childQueue (async MainEventBusProcessor distribution)
-        long startTime = System.currentTimeMillis();
-        long timeout = 2000;
-        while (childQueue.size() == 0 && (System.currentTimeMillis() - startTime) < timeout) {
-            Thread.sleep(50);
-        }
-        assertTrue(childQueue.size() > 0, "Event should arrive in ChildQueue within timeout");
+        // Use callback to wait for event processing instead of polling
+        waitForEventProcessing(() -> mainQueue.enqueueEvent(event));
 
+        // At this point, event has been processed and distributed to childQueue
         childQueue.close(); // Graceful close - events should remain
         assertTrue(childQueue.isClosed());