test_a2a_non_streaming

tests/agent_tests/test_a2a_agent.py

@@ -0,0 +1,71 @@
+"""
+Simple A2A agent tests - non-streaming and streaming.
+
+Requires A2A_AGENT_URL environment variable to be set.
+
+Run with:
+    A2A_AGENT_URL=https://your-agent.example.com pytest tests/agent_tests/test_a2a_agent.py -v -s
+"""
+
+import os
+
+import pytest
+from uuid import uuid4
+
+
+A2A_AGENT_URL = os.environ.get("A2A_AGENT_URL")
+
+
+@pytest.mark.asyncio
+async def test_a2a_non_streaming():
+    """Test non-streaming A2A request."""
+    from a2a.types import MessageSendParams, SendMessageRequest
+    from litellm.a2a_protocol import asend_message
+
+    request = SendMessageRequest(
+        id=str(uuid4()),
+        params=MessageSendParams(
+            message={
+                "role": "user",
+                "parts": [{"kind": "text", "text": "Say hello in one word"}],
+                "messageId": uuid4().hex,
+            }
+        ),
+    )
+
+    response = await asend_message(
+        request=request,
+        api_base=A2A_AGENT_URL,
+    )
+
+    assert response is not None
+    print(f"\nNon-streaming response: {response}")
+
+
+@pytest.mark.asyncio
+async def test_a2a_streaming():
+    """Test streaming A2A request."""
+    from a2a.types import MessageSendParams, SendStreamingMessageRequest
+    from litellm.a2a_protocol import asend_message_streaming
+
+    request = SendStreamingMessageRequest(
+        id=str(uuid4()),
+        params=MessageSendParams(
+            message={
+                "role": "user",
+                "parts": [{"kind": "text", "text": "Say hello in one word"}],
+                "messageId": uuid4().hex,
+            }
+        ),
+    )
+
+    chunks = []
+    async for chunk in asend_message_streaming(
+        request=request,
+        api_base=A2A_AGENT_URL,
+    ):
+        chunks.append(chunk)
+        print(f"\nStreaming chunk: {chunk}")
+
+    assert len(chunks) > 0, "Should receive at least one chunk"
+    print(f"\nTotal chunks received: {len(chunks)}")