fix: fix flaky test testNonBlockingWithMultipleMessages

Author: cmtoan

tests/server-common/src/test/java/io/a2a/server/apps/common/AbstractA2AServerTest.java

@@ -1274,143 +1274,148 @@ public void testDeletePushNotificationConfigSetWithoutConfigId() throws Exceptio
     @Test
     @Timeout(value = 1, unit = TimeUnit.MINUTES)
     public void testNonBlockingWithMultipleMessages() throws Exception {
-        // 1. Send first non-blocking message to create task in WORKING state
-        Message message1 = Message.builder(MESSAGE)
-                .taskId("multi-event-test")
+        String multiEventTaskId = "multi-event-test";
+        try {
+            // 1. Send first non-blocking message to create task in WORKING state
+            Message message1 = Message.builder(MESSAGE)
+                .taskId(multiEventTaskId)
                 .contextId("test-context")
                 .parts(new TextPart("First request"))
                 .build();
 
-        AtomicReference<String> taskIdRef = new AtomicReference<>();
-        CountDownLatch firstTaskLatch = new CountDownLatch(1);
+            AtomicReference<String> taskIdRef = new AtomicReference<>();
+            CountDownLatch firstTaskLatch = new CountDownLatch(1);
 
-        BiConsumer<ClientEvent, AgentCard> firstMessageConsumer = (event, agentCard) -> {
-            if (event instanceof TaskEvent te) {
-                taskIdRef.set(te.getTask().id());
-                firstTaskLatch.countDown();
-            } else if (event instanceof TaskUpdateEvent tue && tue.getUpdateEvent() instanceof TaskStatusUpdateEvent status) {
-                taskIdRef.set(status.taskId());
-                firstTaskLatch.countDown();
-            }
-        };
+            BiConsumer<ClientEvent, AgentCard> firstMessageConsumer = (event, agentCard) -> {
+                if (event instanceof TaskEvent te) {
+                    taskIdRef.set(te.getTask().id());
+                    firstTaskLatch.countDown();
+                } else if (event instanceof TaskUpdateEvent tue && tue.getUpdateEvent() instanceof TaskStatusUpdateEvent status) {
+                    taskIdRef.set(status.taskId());
+                    firstTaskLatch.countDown();
+                }
+            };
 
-        // Non-blocking message creates task in WORKING state and returns immediately
-        // Queue stays open because task is not in final state
-        getPollingClient().sendMessage(message1, List.of(firstMessageConsumer), null);
-
-        assertTrue(firstTaskLatch.await(10, TimeUnit.SECONDS));
-        String taskId = taskIdRef.get();
-        assertNotNull(taskId);
-        assertEquals("multi-event-test", taskId);
-
-        // 2. Resubscribe to task (queue should still be open)
-        CountDownLatch resubEventLatch = new CountDownLatch(2);  // artifact-2 + completion
-        List<io.a2a.spec.UpdateEvent> resubReceivedEvents = new CopyOnWriteArrayList<>();
-        AtomicBoolean resubUnexpectedEvent = new AtomicBoolean(false);
-        AtomicReference<Throwable> resubErrorRef = new AtomicReference<>();
-
-        BiConsumer<ClientEvent, AgentCard> resubConsumer = (event, agentCard) -> {
-            if (event instanceof TaskUpdateEvent tue) {
-                resubReceivedEvents.add(tue.getUpdateEvent());
-                resubEventLatch.countDown();
-            } else {
-                resubUnexpectedEvent.set(true);
-            }
-        };
+            // Non-blocking message creates task in WORKING state and returns immediately
+            // Queue stays open because task is not in final state
+            getPollingClient().sendMessage(message1, List.of(firstMessageConsumer), null);
 
-        Consumer<Throwable> resubErrorHandler = error -> {
-            if (!isStreamClosedError(error)) {
-                resubErrorRef.set(error);
-            }
-        };
+            assertTrue(firstTaskLatch.await(10, TimeUnit.SECONDS));
+            String taskId = taskIdRef.get();
+            assertNotNull(taskId);
+            assertEquals(multiEventTaskId, taskId);
+
+            // 2. Resubscribe to task (queue should still be open)
+            CountDownLatch resubEventLatch = new CountDownLatch(2);  // artifact-2 + completion
+            List<io.a2a.spec.UpdateEvent> resubReceivedEvents = new CopyOnWriteArrayList<>();
+            AtomicBoolean resubUnexpectedEvent = new AtomicBoolean(false);
+            AtomicReference<Throwable> resubErrorRef = new AtomicReference<>();
+
+            BiConsumer<ClientEvent, AgentCard> resubConsumer = (event, agentCard) -> {
+                if (event instanceof TaskUpdateEvent tue) {
+                    resubReceivedEvents.add(tue.getUpdateEvent());
+                    resubEventLatch.countDown();
+                } else {
+                    resubUnexpectedEvent.set(true);
+                }
+            };
 
-        // Wait for subscription to be active
-        CountDownLatch subscriptionLatch = new CountDownLatch(1);
-        awaitStreamingSubscription()
+            Consumer<Throwable> resubErrorHandler = error -> {
+                if (!isStreamClosedError(error)) {
+                    resubErrorRef.set(error);
+                }
+            };
+
+            // Wait for subscription to be active
+            CountDownLatch subscriptionLatch = new CountDownLatch(1);
+            awaitStreamingSubscription()
                 .whenComplete((unused, throwable) -> subscriptionLatch.countDown());
 
-        getClient().resubscribe(new TaskIdParams(taskId),
+            getClient().resubscribe(new TaskIdParams(taskId),
                 List.of(resubConsumer),
                 resubErrorHandler);
 
-        assertTrue(subscriptionLatch.await(15, TimeUnit.SECONDS));
+            assertTrue(subscriptionLatch.await(15, TimeUnit.SECONDS));
 
-        // 3. Send second streaming message to same taskId
-        Message message2 = Message.builder(MESSAGE)
-                .taskId("multi-event-test") // Same taskId
+            // 3. Send second streaming message to same taskId
+            Message message2 = Message.builder(MESSAGE)
+                .taskId(multiEventTaskId) // Same taskId
                 .contextId("test-context")
                 .parts(new TextPart("Second request"))
                 .build();
 
-        CountDownLatch streamEventLatch = new CountDownLatch(2);  // artifact-2 + completion
-        List<io.a2a.spec.UpdateEvent> streamReceivedEvents = new CopyOnWriteArrayList<>();
-        AtomicBoolean streamUnexpectedEvent = new AtomicBoolean(false);
+            CountDownLatch streamEventLatch = new CountDownLatch(2);  // artifact-2 + completion
+            List<io.a2a.spec.UpdateEvent> streamReceivedEvents = new CopyOnWriteArrayList<>();
+            AtomicBoolean streamUnexpectedEvent = new AtomicBoolean(false);
 
-        BiConsumer<ClientEvent, AgentCard> streamConsumer = (event, agentCard) -> {
-            if (event instanceof TaskUpdateEvent tue) {
-                streamReceivedEvents.add(tue.getUpdateEvent());
-                streamEventLatch.countDown();
-            } else {
-                streamUnexpectedEvent.set(true);
-            }
-        };
+            BiConsumer<ClientEvent, AgentCard> streamConsumer = (event, agentCard) -> {
+                if (event instanceof TaskUpdateEvent tue) {
+                    streamReceivedEvents.add(tue.getUpdateEvent());
+                    streamEventLatch.countDown();
+                } else {
+                    streamUnexpectedEvent.set(true);
+                }
+            };
 
-        // Streaming message adds artifact-2 and completes task
-        getClient().sendMessage(message2, List.of(streamConsumer), null);
+            // Streaming message adds artifact-2 and completes task
+            getClient().sendMessage(message2, List.of(streamConsumer), null);
 
-        // 4. Verify both consumers received artifact-2 and completion
-        assertTrue(resubEventLatch.await(10, TimeUnit.SECONDS));
-        assertTrue(streamEventLatch.await(10, TimeUnit.SECONDS));
+            // 4. Verify both consumers received artifact-2 and completion
+            assertTrue(resubEventLatch.await(10, TimeUnit.SECONDS));
+            assertTrue(streamEventLatch.await(10, TimeUnit.SECONDS));
 
-        assertFalse(resubUnexpectedEvent.get());
-        assertFalse(streamUnexpectedEvent.get());
-        assertNull(resubErrorRef.get());
+            assertFalse(resubUnexpectedEvent.get());
+            assertFalse(streamUnexpectedEvent.get());
+            assertNull(resubErrorRef.get());
 
-        // Both should have received 2 events: artifact-2 and completion
-        assertEquals(2, resubReceivedEvents.size());
-        assertEquals(2, streamReceivedEvents.size());
+            // Both should have received 2 events: artifact-2 and completion
+            assertEquals(2, resubReceivedEvents.size());
+            assertEquals(2, streamReceivedEvents.size());
 
-        // Verify resubscription events
-        long resubArtifactCount = resubReceivedEvents.stream()
+            // Verify resubscription events
+            long resubArtifactCount = resubReceivedEvents.stream()
                 .filter(e -> e instanceof TaskArtifactUpdateEvent)
                 .count();
-        assertEquals(1, resubArtifactCount);
+            assertEquals(1, resubArtifactCount);
 
-        long resubCompletionCount = resubReceivedEvents.stream()
+            long resubCompletionCount = resubReceivedEvents.stream()
                 .filter(e -> e instanceof TaskStatusUpdateEvent)
                 .filter(e -> ((TaskStatusUpdateEvent) e).isFinal())
                 .count();
-        assertEquals(1, resubCompletionCount);
+            assertEquals(1, resubCompletionCount);
 
-        // Verify streaming events
-        long streamArtifactCount = streamReceivedEvents.stream()
+            // Verify streaming events
+            long streamArtifactCount = streamReceivedEvents.stream()
                 .filter(e -> e instanceof TaskArtifactUpdateEvent)
                 .count();
-        assertEquals(1, streamArtifactCount);
+            assertEquals(1, streamArtifactCount);
 
-        long streamCompletionCount = streamReceivedEvents.stream()
+            long streamCompletionCount = streamReceivedEvents.stream()
                 .filter(e -> e instanceof TaskStatusUpdateEvent)
                 .filter(e -> ((TaskStatusUpdateEvent) e).isFinal())
                 .count();
-        assertEquals(1, streamCompletionCount);
+            assertEquals(1, streamCompletionCount);
 
-        // Verify artifact-2 details from resubscription
-        TaskArtifactUpdateEvent resubArtifact = (TaskArtifactUpdateEvent) resubReceivedEvents.stream()
+            // Verify artifact-2 details from resubscription
+            TaskArtifactUpdateEvent resubArtifact = (TaskArtifactUpdateEvent) resubReceivedEvents.stream()
                 .filter(e -> e instanceof TaskArtifactUpdateEvent)
                 .findFirst()
                 .orElseThrow();
-        assertEquals("artifact-2", resubArtifact.artifact().artifactId());
-        assertEquals("Second message artifact",
+            assertEquals("artifact-2", resubArtifact.artifact().artifactId());
+            assertEquals("Second message artifact",
                 ((TextPart) resubArtifact.artifact().parts().get(0)).text());
 
-        // Verify artifact-2 details from streaming
-        TaskArtifactUpdateEvent streamArtifact = (TaskArtifactUpdateEvent) streamReceivedEvents.stream()
+            // Verify artifact-2 details from streaming
+            TaskArtifactUpdateEvent streamArtifact = (TaskArtifactUpdateEvent) streamReceivedEvents.stream()
                 .filter(e -> e instanceof TaskArtifactUpdateEvent)
                 .findFirst()
                 .orElseThrow();
-        assertEquals("artifact-2", streamArtifact.artifact().artifactId());
-        assertEquals("Second message artifact",
+            assertEquals("artifact-2", streamArtifact.artifact().artifactId());
+            assertEquals("Second message artifact",
                 ((TextPart) streamArtifact.artifact().parts().get(0)).text());
+        } finally {
+            deleteTaskInTaskStore(multiEventTaskId);
+        }
     }
 
     @Test