test: loop until result on A2A tests (#3233)

Author: DouweM

tests/test_a2a.py

@@ -560,11 +560,11 @@ def track_messages(messages: list[ModelMessage], info: AgentInfo) -> ModelRespon
             task1_id = result1['id']
             context_id = result1['context_id']
 
-            # Wait for first task to complete
-            await anyio.sleep(0.1)
-            task1 = await a2a_client.get_task(task1_id)
-            assert 'result' in task1
-            assert task1['result']['status']['state'] == 'completed'
+            while task1 := await a2a_client.get_task(task1_id):  # pragma: no branch
+                if 'result' in task1 and task1['result']['status']['state'] == 'completed':
+                    result1 = task1['result']
+                    break
+                await anyio.sleep(0.1)
 
             # Verify the model received at least one message
             assert len(messages_received) == 1
@@ -668,11 +668,13 @@ def return_thinking_response(_: list[ModelMessage], info: AgentInfo) -> ModelRes
             task_id = result['id']
 
             # Wait for completion
-            await anyio.sleep(0.1)
-            task = await a2a_client.get_task(task_id)
+            while task := await a2a_client.get_task(task_id):  # pragma: no branch
+                if 'result' in task and task['result']['status']['state'] == 'completed':
+                    result = task['result']
+                    break
+                await anyio.sleep(0.1)
 
-            assert 'result' in task
-            assert task['result'] == snapshot(
+            assert result == snapshot(
                 {
                     'id': IsStr(),
                     'context_id': IsStr(),