Gemini feedback and fixing tests

Author: kabir

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

@@ -357,7 +357,10 @@ void testReplicatedEventProcessedWhenTaskActive() throws InterruptedException {
         EventQueue childQueue = queue.tap();
         Event dequeuedEvent;
         try {
-            dequeuedEvent = childQueue.dequeueEventItem(100).getEvent();
+            // 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;

server-common/src/main/java/io/a2a/server/events/EventQueue.java

@@ -306,21 +306,19 @@ void childClosing(ChildQueue child, boolean immediate) {
          * Called by MainEventBusProcessor after TaskStore persistence.
          */
         void distributeToChildren(EventQueueItem item) {
-            synchronized (children) {
-                int childCount = children.size();
-                if (LOGGER.isDebugEnabled()) {
-                    LOGGER.debug("MainQueue[{}]: Distributing event {} to {} children",
-                                taskId, item.getEvent().getClass().getSimpleName(), childCount);
-                }
-                children.forEach(child -> {
-                    LOGGER.debug("MainQueue[{}]: Enqueueing event {} to child queue",
-                                taskId, item.getEvent().getClass().getSimpleName());
-                    child.internalEnqueueItem(item);
-                });
-                if (LOGGER.isDebugEnabled()) {
-                    LOGGER.debug("MainQueue[{}]: Completed distribution of {} to {} children",
-                                taskId, item.getEvent().getClass().getSimpleName(), childCount);
-                }
+            int childCount = children.size();
+            if (LOGGER.isDebugEnabled()) {
+                LOGGER.debug("MainQueue[{}]: Distributing event {} to {} children",
+                            taskId, item.getEvent().getClass().getSimpleName(), childCount);
+            }
+            children.forEach(child -> {
+                LOGGER.debug("MainQueue[{}]: Enqueueing event {} to child queue",
+                            taskId, item.getEvent().getClass().getSimpleName());
+                child.internalEnqueueItem(item);
+            });
+            if (LOGGER.isDebugEnabled()) {
+                LOGGER.debug("MainQueue[{}]: Completed distribution of {} to {} children",
+                            taskId, item.getEvent().getClass().getSimpleName(), childCount);
             }
         }
 
@@ -471,8 +469,11 @@ protected void doClose(boolean immediate) {
             super.doClose(immediate);  // Sets closed flag
             if (immediate) {
                 // Immediate close: clear pending events from local queue
+                int clearedCount = queue.size();
                 queue.clear();
-                LOGGER.debug("Cleared ChildQueue for immediate close: {}", this);
+                // Release semaphore permits for cleared events to prevent deadlock
+                semaphore.release(clearedCount);
+                LOGGER.debug("Cleared {} events from ChildQueue for immediate close: {}", clearedCount, this);
             }
             // For graceful close, let the queue drain naturally through normal consumption
         }

server-common/src/main/java/io/a2a/server/events/MainEventBusProcessor.java

@@ -44,7 +44,7 @@ public class MainEventBusProcessor implements Runnable {
      * Default is NOOP to avoid null checks in production code.
      * Tests can inject their own callback via setCallback().
      */
-    private static volatile MainEventBusProcessorCallback callback = MainEventBusProcessorCallback.NOOP;
+    private volatile MainEventBusProcessorCallback callback = MainEventBusProcessorCallback.NOOP;
 
     private final MainEventBus eventBus;
 
@@ -71,14 +71,14 @@ public MainEventBusProcessor(MainEventBus eventBus, TaskStore taskStore, PushNot
      *
      * @param callback the callback to invoke during event processing, or null for NOOP
      */
-    public static void setCallback(MainEventBusProcessorCallback callback) {
-        MainEventBusProcessor.callback = callback != null ? callback : MainEventBusProcessorCallback.NOOP;
+    public void setCallback(MainEventBusProcessorCallback callback) {
+        this.callback = callback != null ? callback : MainEventBusProcessorCallback.NOOP;
     }
 
     @PostConstruct
     void start() {
         processorThread = new Thread(this, "MainEventBusProcessor");
-        processorThread.setDaemon(false); // Keep JVM alive
+        processorThread.setDaemon(true); // Allow JVM to exit even if this thread is running
         processorThread.start();
         LOGGER.info("MainEventBusProcessor started");
     }
@@ -98,9 +98,13 @@ void stop() {
         if (processorThread != null) {
             processorThread.interrupt();
             try {
+                long start = System.currentTimeMillis();
                 processorThread.join(5000); // Wait up to 5 seconds
+                long elapsed = System.currentTimeMillis() - start;
+                LOGGER.info("MainEventBusProcessor thread stopped in {}ms", elapsed);
             } catch (InterruptedException e) {
                 Thread.currentThread().interrupt();
+                LOGGER.warn("Interrupted while waiting for MainEventBusProcessor thread to stop");
             }
         }
         LOGGER.info("MainEventBusProcessor stopped");

server-common/src/main/java/io/a2a/server/events/MainEventBusProcessorCallback.java

@@ -12,10 +12,13 @@
  * Usage in tests:
  * <pre>
  * {@code
+ * @Inject
+ * MainEventBusProcessor processor;
+ *
  * @BeforeEach
  * void setUp() {
  *     CountDownLatch latch = new CountDownLatch(3);
- *     MainEventBusProcessor.setCallback(new MainEventBusProcessorCallback() {
+ *     processor.setCallback(new MainEventBusProcessorCallback() {
  *         public void onEventProcessed(String taskId, Event event) {
  *             latch.countDown();
  *         }
@@ -24,7 +27,7 @@
  *
  * @AfterEach
  * void tearDown() {
- *     MainEventBusProcessor.setCallback(null); // Reset to NOOP
+ *     processor.setCallback(null); // Reset to NOOP
  * }
  * }
  * </pre>

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

@@ -15,6 +15,8 @@
 import java.util.concurrent.atomic.AtomicReference;
 
 import com.fasterxml.jackson.core.JsonProcessingException;
+import io.a2a.server.tasks.InMemoryTaskStore;
+import io.a2a.server.tasks.PushNotificationSender;
 import io.a2a.spec.A2AError;
 import io.a2a.spec.A2AServerException;
 import io.a2a.spec.Artifact;
@@ -28,14 +30,18 @@
 import io.a2a.spec.TaskStatusUpdateEvent;
 import io.a2a.spec.TextPart;
 import io.a2a.util.Utils;
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
 public class EventConsumerTest {
 
+    private static final PushNotificationSender NOOP_PUSHNOTIFICATION_SENDER = task -> {};
+
     private EventQueue eventQueue;
     private EventConsumer eventConsumer;
-
+    private MainEventBus mainEventBus;
+    private MainEventBusProcessor mainEventBusProcessor;
 
     private static final String MINIMAL_TASK = """
             {
@@ -57,10 +63,25 @@ public class EventConsumerTest {
 
     @BeforeEach
     public void init() {
-        eventQueue = EventQueueUtil.getEventQueueBuilder().build().tap();
+        // Set up MainEventBus and processor for production-like test environment
+        InMemoryTaskStore taskStore = new InMemoryTaskStore();
+        mainEventBus = new MainEventBus();
+        mainEventBusProcessor = new MainEventBusProcessor(mainEventBus, taskStore, NOOP_PUSHNOTIFICATION_SENDER);
+        EventQueueUtil.start(mainEventBusProcessor);
+
+        eventQueue = EventQueueUtil.getEventQueueBuilder()
+                .mainEventBus(mainEventBus)
+                .build().tap();
         eventConsumer = new EventConsumer(eventQueue);
     }
 
+    @AfterEach
+    public void cleanup() {
+        if (mainEventBusProcessor != null) {
+            EventQueueUtil.stop(mainEventBusProcessor);
+        }
+    }
+
     @Test
     public void testConsumeOneTaskEvent() throws Exception {
         Task event = Utils.unmarshalFrom(MINIMAL_TASK, Task.TYPE_REFERENCE);
@@ -343,7 +364,9 @@ public void onComplete() {
 
     @Test
     public void testConsumeAllStopsOnQueueClosed() throws Exception {
-        EventQueue queue = EventQueueUtil.getEventQueueBuilder().build().tap();
+        EventQueue queue = EventQueueUtil.getEventQueueBuilder()
+                .mainEventBus(mainEventBus)
+                .build().tap();
         EventConsumer consumer = new EventConsumer(queue);
 
         // Close the queue immediately
@@ -389,7 +412,9 @@ public void onComplete() {
 
     @Test
     public void testConsumeAllHandlesQueueClosedException() throws Exception {
-        EventQueue queue = EventQueueUtil.getEventQueueBuilder().build().tap();
+        EventQueue queue = EventQueueUtil.getEventQueueBuilder()
+                .mainEventBus(mainEventBus)
+                .build().tap();
         EventConsumer consumer = new EventConsumer(queue);
 
         // Add a message event (which will complete the stream)
@@ -447,7 +472,9 @@ public void onComplete() {
 
     @Test
     public void testConsumeAllTerminatesOnQueueClosedEvent() throws Exception {
-        EventQueue queue = EventQueueUtil.getEventQueueBuilder().build().tap();
+        EventQueue queue = EventQueueUtil.getEventQueueBuilder()
+                .mainEventBus(mainEventBus)
+                .build().tap();
         EventConsumer consumer = new EventConsumer(queue);
 
         // Enqueue a QueueClosedEvent (poison pill)

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

@@ -225,7 +225,9 @@ public void testCloseImmediatePropagationToChildren() throws Exception {
 
     @Test
     public void testEnqueueEventWhenClosed() throws Exception {
-        EventQueue mainQueue = EventQueueUtil.getEventQueueBuilder().build();
+        EventQueue mainQueue = EventQueueUtil.getEventQueueBuilder()
+                .mainEventBus(mainEventBus)
+                .build();
         EventQueue childQueue = mainQueue.tap();
         Event event = Utils.unmarshalFrom(MINIMAL_TASK, Task.TYPE_REFERENCE);
 
@@ -260,16 +262,22 @@ public void testDequeueEventWhenClosedAndEmpty() throws Exception {
 
     @Test
     public void testDequeueEventWhenClosedButHasEvents() throws Exception {
-        EventQueue mainQueue = EventQueueUtil.getEventQueueBuilder().build();
+        EventQueue mainQueue = EventQueueUtil.getEventQueueBuilder()
+                .mainEventBus(mainEventBus)
+                .build();
         EventQueue childQueue = mainQueue.tap();
         Event event = Utils.unmarshalFrom(MINIMAL_TASK, Task.TYPE_REFERENCE);
 
         // Enqueue to mainQueue
         mainQueue.enqueueEvent(event);
 
-        // Wait for event to arrive in childQueue (use peek-like behavior by dequeueing then re-checking)
-        // Actually, just wait a bit for async processing
-        Thread.sleep(100); // Give MainEventBusProcessor time to distribute 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");
 
         childQueue.close(); // Graceful close - events should remain
         assertTrue(childQueue.isClosed());

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

@@ -1,102 +1,43 @@
 package io.a2a.server.events;
 
-import io.a2a.server.tasks.InMemoryTaskStore;
-import io.a2a.server.tasks.PushNotificationSender;
-import io.a2a.server.tasks.TaskStateProvider;
 import java.util.concurrent.atomic.AtomicInteger;
 
 public class EventQueueUtil {
     // Shared MainEventBus for all tests (to avoid creating one per test)
+    // Each test creates its own MainEventBusProcessor in @BeforeEach/@AfterEach for proper isolation
     private static final MainEventBus TEST_EVENT_BUS = new MainEventBus();
 
-    // Shared MainEventBusProcessor for all tests (automatically processes events)
-    private static final MainEventBusProcessor TEST_PROCESSOR;
-
-    static {
-        // Initialize and start the processor once for all tests
-        InMemoryTaskStore testTaskStore = new InMemoryTaskStore();
-        PushNotificationSender testPushSender = taskId -> {}; // No-op for tests
-        TEST_PROCESSOR = new MainEventBusProcessor(TEST_EVENT_BUS, testTaskStore, testPushSender);
-        TEST_PROCESSOR.start();  // Start background thread
-
-        // Register shutdown hook to stop processor
-        Runtime.getRuntime().addShutdownHook(new Thread(() -> TEST_PROCESSOR.stop()));
-    }
-
     // Counter for generating unique test taskIds
     private static final AtomicInteger TASK_ID_COUNTER = new AtomicInteger(0);
 
-    // Since EventQueue.builder() is package protected, add a method to expose it
-    // Note: Now includes MainEventBus requirement and default taskId
-    // Returns MainQueue - tests should call .tap() if they need to consume events
+    /**
+     * Get an EventQueue builder pre-configured with the shared test MainEventBus and a unique taskId.
+     * <p>
+     * Note: Returns MainQueue - tests should call .tap() if they need to consume events.
+     * </p>
+     *
+     * @return builder with TEST_EVENT_BUS and unique taskId already set
+     */
     public static EventQueue.EventQueueBuilder getEventQueueBuilder() {
-        return new EventQueueBuilderWrapper(
-                EventQueue.builder(TEST_EVENT_BUS)
-                        .taskId("test-task-" + TASK_ID_COUNTER.incrementAndGet())
-        );
-    }
-
-    // Get the shared test MainEventBus instance
-    public static MainEventBus getTestEventBus() {
-        return TEST_EVENT_BUS;
+        return EventQueue.builder(TEST_EVENT_BUS)
+                .taskId("test-task-" + TASK_ID_COUNTER.incrementAndGet());
     }
 
+    /**
+     * Start a MainEventBusProcessor instance.
+     *
+     * @param processor the processor to start
+     */
     public static void start(MainEventBusProcessor processor) {
         processor.start();
     }
 
+    /**
+     * Stop a MainEventBusProcessor instance.
+     *
+     * @param processor the processor to stop
+     */
     public static void stop(MainEventBusProcessor processor) {
         processor.stop();
     }
-
-    // Wrapper that delegates to actual builder
-    private static class EventQueueBuilderWrapper extends EventQueue.EventQueueBuilder {
-        private final EventQueue.EventQueueBuilder delegate;
-
-        EventQueueBuilderWrapper(EventQueue.EventQueueBuilder delegate) {
-            this.delegate = delegate;
-        }
-
-        @Override
-        public EventQueue.EventQueueBuilder queueSize(int queueSize) {
-            delegate.queueSize(queueSize);
-            return this;
-        }
-
-        @Override
-        public EventQueue.EventQueueBuilder hook(EventEnqueueHook hook) {
-            delegate.hook(hook);
-            return this;
-        }
-
-        @Override
-        public EventQueue.EventQueueBuilder taskId(String taskId) {
-            delegate.taskId(taskId);
-            return this;
-        }
-
-        @Override
-        public EventQueue.EventQueueBuilder addOnCloseCallback(Runnable onCloseCallback) {
-            delegate.addOnCloseCallback(onCloseCallback);
-            return this;
-        }
-
-        @Override
-        public EventQueue.EventQueueBuilder taskStateProvider(TaskStateProvider taskStateProvider) {
-            delegate.taskStateProvider(taskStateProvider);
-            return this;
-        }
-
-        @Override
-        public EventQueue.EventQueueBuilder mainEventBus(MainEventBus mainEventBus) {
-            delegate.mainEventBus(mainEventBus);
-            return this;
-        }
-
-        @Override
-        public EventQueue build() {
-            // Return MainQueue directly - tests should call .tap() if they need ChildQueue
-            return delegate.build();
-        }
-    }
 }