chore: tighten up tests (#300)

Author: kabir

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

@@ -9,7 +9,9 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.Flow;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
 
 import com.fasterxml.jackson.core.JsonProcessingException;
@@ -297,13 +299,16 @@ public void run() {
     }
 
     @Test
-    public void testConsumeAllRaisesStoredException() {
+    public void testConsumeAllRaisesStoredException() throws InterruptedException {
         // Set an error in the event consumer
         setEventConsumerError(new RuntimeException("Stored error"));
 
         Flow.Publisher<Event> publisher = eventConsumer.consumeAll();
         final AtomicReference<Throwable> receivedError = new AtomicReference<>();
 
+        final CountDownLatch errorLatch = new CountDownLatch(1);
+
+
         publisher.subscribe(new Flow.Subscriber<>() {
             @Override
             public void onSubscribe(Flow.Subscription subscription) {
@@ -313,19 +318,25 @@ public void onSubscribe(Flow.Subscription subscription) {
             @Override
             public void onNext(Event item) {
                 // Should not be called
+                errorLatch.countDown();
             }
 
             @Override
             public void onError(Throwable throwable) {
                 receivedError.set(throwable);
+                errorLatch.countDown();
             }
 
             @Override
             public void onComplete() {
                 // Should not be called
+                errorLatch.countDown();
             }
         });
 
+        // Wait for error callback with timeout
+        assertTrue(errorLatch.await(5, TimeUnit.SECONDS), "Test timed out waiting for onError callback.");
+
         assertNotNull(receivedError.get());
         assertEquals("Stored error", receivedError.get().getMessage());
     }
@@ -341,6 +352,7 @@ public void testConsumeAllStopsOnQueueClosed() throws Exception {
         Flow.Publisher<Event> publisher = consumer.consumeAll();
         final List<Event> receivedEvents = new ArrayList<>();
         final AtomicReference<Boolean> completed = new AtomicReference<>(false);
+        final CountDownLatch completionLatch = new CountDownLatch(1);
 
         publisher.subscribe(new Flow.Subscriber<>() {
             @Override
@@ -356,14 +368,19 @@ public void onNext(Event item) {
             @Override
             public void onError(Throwable throwable) {
                 // Should not be called
+                completionLatch.countDown();
             }
 
             @Override
             public void onComplete() {
                 completed.set(true);
+                completionLatch.countDown();
             }
         });
 
+        // Wait for completion with timeout
+        assertTrue(completionLatch.await(5, TimeUnit.SECONDS), "Test timed out waiting for onComplete callback.");
+
         // Should complete immediately with no events
         assertTrue(completed.get());
         assertEquals(0, receivedEvents.size());
@@ -385,6 +402,7 @@ public void testConsumeAllHandlesQueueClosedException() throws Exception {
         Flow.Publisher<Event> publisher = consumer.consumeAll();
         final List<Event> receivedEvents = new ArrayList<>();
         final AtomicReference<Boolean> completed = new AtomicReference<>(false);
+        final CountDownLatch completionLatch = new CountDownLatch(1);
 
         publisher.subscribe(new Flow.Subscriber<>() {
             @Override
@@ -400,14 +418,19 @@ public void onNext(Event item) {
             @Override
             public void onError(Throwable throwable) {
                 // Should not be called
+                completionLatch.countDown();
             }
 
             @Override
             public void onComplete() {
                 completed.set(true);
+                completionLatch.countDown();
             }
         });
 
+        // Wait for completion with timeout
+        assertTrue(completionLatch.await(5, TimeUnit.SECONDS), "Test timed out waiting for onComplete callback.");
+
         // Should have received the message and completed
         assertTrue(completed.get());
         assertEquals(1, receivedEvents.size());
@@ -426,7 +449,7 @@ private Throwable getEventConsumerError() {
             java.lang.reflect.Field errorField = EventConsumer.class.getDeclaredField("error");
             errorField.setAccessible(true);
             return (Throwable) errorField.get(eventConsumer);
-        } catch (Exception e) {
+        } catch (NoSuchFieldException | IllegalAccessException e) {
             throw new RuntimeException("Failed to access error field", e);
         }
     }
@@ -436,7 +459,7 @@ private void setEventConsumerError(Throwable error) {
             java.lang.reflect.Field errorField = EventConsumer.class.getDeclaredField("error");
             errorField.setAccessible(true);
             errorField.set(eventConsumer, error);
-        } catch (Exception e) {
+        } catch (NoSuchFieldException | IllegalAccessException e) {
             throw new RuntimeException("Failed to set error field", e);
         }
     }

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

@@ -183,12 +183,8 @@ public void testCloseImmediatePropagationToChildren() throws Exception {
         assertTrue(childQueue.isClosed());
 
         // Child queue should be cleared due to immediate close
-        try {
-            Event result = childQueue.dequeueEvent(-1);
-            assertNull(result);
-        } catch (EventQueueClosedException e) {
-            // This is expected when queue is closed
-        }
+        // Child queue should be cleared and closed, so dequeueing should throw
+        assertThrows(EventQueueClosedException.class, () -> childQueue.dequeueEvent(-1));
     }
 
     @Test
@@ -203,12 +199,8 @@ public void testEnqueueEventWhenClosed() throws Exception {
         queue.enqueueEvent(event);
 
         // Verify the queue is still empty
-        try {
-            Event dequeuedEvent = queue.dequeueEvent(-1);
-            assertNull(dequeuedEvent);
-        } catch (EventQueueClosedException e) {
-            // This is expected when queue is closed and empty
-        }
+        // Verify the queue is still empty and dequeueing from a closed, empty queue throws
+        assertThrows(EventQueueClosedException.class, () -> queue.dequeueEvent(-1));
     }
 
     @Test

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

@@ -76,6 +76,9 @@ public void testTapExistingQueue() {
         EventQueue tappedQueue = queueManager.tap(taskId);
 
         assertNotNull(tappedQueue);
+
+        // Tapped queue should be different from original
+        // (it's a child queue in Java implementation)
         assertNotSame(queue, tappedQueue, "Tapped queue should be a different instance from the original.");
     }
 
@@ -205,9 +208,11 @@ public void testCreateOrTapRaceCondition() throws InterruptedException, Executio
         EventQueue managerQueue = queueManager.get(taskId);
         assertNotNull(managerQueue);
 
-        // One of the results should be the main queue instance, and the rest should be
-        // different (tapped) queue instances.
+        // One of the results should be the main queue instance, the others should be tapped child queues.
         long mainQueueCount = results.stream().filter(q -> q == managerQueue).count();
-        assertEquals(1, mainQueueCount, "Exactly one returned queue should be the main queue instance.");
+        assertEquals(1, mainQueueCount, "Exactly one thread should have created the main queue.");
+
+        long childQueueCount = results.stream().filter(q -> q != managerQueue).count();
+        assertEquals(results.size() - 1, childQueueCount, "All other threads should have received a tapped queue.");
     }
 }