update test_client.py

Author: wadave

examples/a2a-adk-app/airbnb_agent/test_client.py

@@ -5,34 +5,36 @@
 
 from a2a.client import A2AClient
 from a2a.types import (
-    GetTaskResponse,
     SendMessageResponse,
+    GetTaskResponse,
     SendMessageSuccessResponse,
     Task,
-    TaskState,
+    SendMessageRequest,
+    MessageSendParams,
+    GetTaskRequest,
+    TaskQueryParams,
 )
-
+import traceback
 
 AGENT_URL = "http://localhost:10002"
 
-
 def create_send_message_payload(
     text: str, task_id: str | None = None, context_id: str | None = None
 ) -> dict[str, Any]:
     """Helper function to create the payload for sending a task."""
     payload: dict[str, Any] = {
-        "message": {
-            "role": "user",
-            "parts": [{"type": "text", "text": text}],
-            "messageId": uuid4().hex,
+        'message': {
+            'role': 'user',
+            'parts': [{'kind': 'text', 'text': text}],
+            'messageId': uuid4().hex,
         },
     }
 
     if task_id:
-        payload["message"]["taskId"] = task_id
+        payload['message']['taskId'] = task_id
 
     if context_id:
-        payload["message"]["contextId"] = context_id
+        payload['message']['contextId'] = context_id
     return payload
 
 
@@ -49,98 +51,46 @@ async def run_single_turn_test(client: A2AClient) -> None:
     """Runs a single-turn non-streaming test."""
 
     send_payload = create_send_message_payload(
-        text="Please find a bedroom in new york city, June 15-20, 2025"
+        text='Please find a bedroom in new york city, June 15-20, 2025'
     )
-    print("send_payload", send_payload)
+    request = SendMessageRequest(params=MessageSendParams(**send_payload))
+
+    print('--- Single Turn Request ---')
     # Send Message
-    send_response: SendMessageResponse = await client.send_message(payload=send_payload)
-    print("send_response", send_response)
-    print_json_response(send_response, "Single Turn Request Response")
+    send_response: SendMessageResponse = await client.send_message(request)
+    print_json_response(send_response, 'Single Turn Request Response')
     if not isinstance(send_response.root, SendMessageSuccessResponse):
-        print("received non-success response. Aborting get task ")
+        print('received non-success response. Aborting get task ')
         return
 
     if not isinstance(send_response.root.result, Task):
-        print("received non-task response. Aborting get task ")
+        print('received non-task response. Aborting get task ')
         return
 
     task_id: str = send_response.root.result.id
-    print("---Query Task---")
+    print('---Query Task---')
     # query the task
-    task_id_payload = {"id": task_id}
-    get_response: GetTaskResponse = await client.get_task(payload=task_id_payload)
-    print_json_response(get_response, "Query Task Response")
-
-
-async def run_streaming_test(client: A2AClient) -> None:
-    """Runs a single-turn streaming test."""
-
-    send_payload = create_send_message_payload(
-        text="Please find a room in LA, CA, june 15, 2025, checkout date is june 18, 2 adults"
-    )
-
-    print("--- Single Turn Streaming Request ---")
-    stream_response = client.send_message_streaming(payload=send_payload)
-    async for chunk in stream_response:
-        print_json_response(chunk, "Streaming Chunk")
-
-
-async def run_multi_turn_test(client: A2AClient) -> None:
-    """Runs a multi-turn non-streaming test."""
-    print("--- Multi-Turn Request ---")
-    # --- First Turn ---
-
-    first_turn_payload = create_send_message_payload(
-        text="Please find a room in LA, CA, june 15, 2025, checkout date is june 18, 2 adults"
-    )
-    first_turn_response: SendMessageResponse = await client.send_message(
-        payload=first_turn_payload
-    )
-    print_json_response(first_turn_response, "Multi-Turn: First Turn Response")
-
-    context_id: str | None = None
-    if isinstance(first_turn_response.root, SendMessageSuccessResponse) and isinstance(
-        first_turn_response.root.result, Task
-    ):
-        task: Task = first_turn_response.root.result
-        context_id = task.contextId  # Capture context ID
-
-        # --- Second Turn (if input required) ---
-        if task.status.state == TaskState.input_required and context_id:
-            print("--- Multi-Turn: Second Turn (Input Required) ---")
-            second_turn_payload = create_send_message_payload(
-                "in NYC", task.id, context_id
-            )
-            second_turn_response = await client.send_message(
-                payload=second_turn_payload
-            )
-            print_json_response(
-                second_turn_response, "Multi-Turn: Second Turn Response"
-            )
-        elif not context_id:
-            print("Warning: Could not get context ID from first turn response.")
-        else:
-            print("First turn completed, no further input required for this test case.")
-
+    get_request = GetTaskRequest(params=TaskQueryParams(id=task_id))
+    get_response: GetTaskResponse = await client.get_task(get_request)
+    print_json_response(get_response, 'Query Task Response')
 
+    
 async def main() -> None:
     """Main function to run the tests."""
-    print(f"Connecting to agent at {AGENT_URL}...")
+    print(f'Connecting to agent at {AGENT_URL}...')
     try:
         async with httpx.AsyncClient(timeout=30) as httpx_client:
             client = await A2AClient.get_client_from_agent_card_url(
                 httpx_client, AGENT_URL
             )
-            print("Connection successful.")
+            print('Connection successful.')
 
             await run_single_turn_test(client)
-            # await run_streaming_test(client)
-            # await run_multi_turn_test(client)
 
     except Exception as e:
-        print(f"An error occurred: {e}")
-        print("Ensure the agent server is running.")
-
+        traceback.print_exc()
+        print(f'An error occurred: {e}')
+        print('Ensure the agent server is running.')
 
 if __name__ == "__main__":
     import asyncio