fix: Eliminate race condition in testInputRequiredWorkflow

kabir · kabir · commit b0a71242b8e0 · 2026-02-05T16:49:01.000Z
The test consumers were counting down their latches on the FIRST TaskEvent
received, which could be the intermediate WORKING state instead of the
expected terminal state (INPUT_REQUIRED or COMPLETED).

This caused intermittent test failures with:
  expected: &lt;INPUT_REQUIRED&gt; but was: &lt;WORKING&gt;

Fix: Both consumers now only count down when receiving the expected
terminal state:
- initialConsumer waits for INPUT_REQUIRED (not WORKING)
- completionConsumer waits for COMPLETED (not WORKING)

This matches the agent's event emission pattern:
  agentEmitter.startWork()    // WORKING state
  agentEmitter.requiresInput(...) // INPUT_REQUIRED state
  ...later...
  agentEmitter.startWork()    // WORKING state
  agentEmitter.complete()     // COMPLETED state
diff --git a/tests/server-common/src/test/java/io/a2a/server/apps/common/AbstractA2AServerTest.java b/tests/server-common/src/test/java/io/a2a/server/apps/common/AbstractA2AServerTest.java
@@ -1513,8 +1513,12 @@ public void testInputRequiredWorkflow() throws Exception {
 
             BiConsumer<ClientEvent, AgentCard> initialConsumer = (event, agentCard) -> {
                 if (event instanceof TaskEvent te) {
-                    initialState.set(te.getTask().status().state());
-                    initialLatch.countDown();
+                    TaskState state = te.getTask().status().state();
+                    initialState.set(state);
+                    // Only count down when we receive INPUT_REQUIRED, not intermediate states like WORKING
+                    if (state == TaskState.INPUT_REQUIRED) {
+                        initialLatch.countDown();
+                    }
                 } else {
                     initialUnexpectedEvent.set(true);
                 }
@@ -1539,8 +1543,12 @@ public void testInputRequiredWorkflow() throws Exception {
 
             BiConsumer<ClientEvent, AgentCard> completionConsumer = (event, agentCard) -> {
                 if (event instanceof TaskEvent te) {
-                    completedState.set(te.getTask().status().state());
-                    completionLatch.countDown();
+                    TaskState state = te.getTask().status().state();
+                    completedState.set(state);
+                    // Only count down when we receive COMPLETED, not intermediate states like WORKING
+                    if (state == TaskState.COMPLETED) {
+                        completionLatch.countDown();
+                    }
                 } else {
                     completionUnexpectedEvent.set(true);
                 }
