diff --git a/tests/agent_tests/test_a2a_agent.py b/tests/agent_tests/test_a2a_agent.py new file mode 100644 index 00000000000..131f0d1cd39 --- /dev/null +++ b/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)}")