From 14906a1777f057d9f934b9747ef26fd8842e8426 Mon Sep 17 00:00:00 2001 From: Faheem Shah Umer Vattam Kandathil <161448038+faheem-shah-umer@users.noreply.github.com> Date: Thu, 22 Oct 2026 23:02:00 +0200 Subject: [PATCH 1/2] test: extend selected scenarios to REST --- tests/integration/test_scenario_cleanup.py | 218 ++++++++++++++++++++ tests/integration/test_scenarios.py | 220 ++++++++++++++++----- 2 files changed, 392 insertions(+), 46 deletions(-) create mode 100644 tests/integration/test_scenario_cleanup.py diff --git a/tests/integration/test_scenario_cleanup.py b/tests/integration/test_scenario_cleanup.py new file mode 100644 index 000000000..66a8b51f3 --- /dev/null +++ b/tests/integration/test_scenario_cleanup.py @@ -0,0 +1,218 @@ +import asyncio +import socket + +from collections.abc import AsyncGenerator, AsyncIterator +from contextlib import AsyncExitStack + +import grpc +import httpx +import pytest +import pytest_asyncio + +from a2a.client.client import Client +from a2a.client.client_factory import ClientFactory +from a2a.helpers.proto_helpers import new_task_from_user_message +from a2a.server.agent_execution import AgentExecutor, RequestContext +from a2a.server.events import EventQueue +from a2a.types.a2a_pb2 import ( + AgentCard, + Message, + Part, + Role, + SendMessageRequest, + StreamResponse, + TaskState, +) +from a2a.utils import TransportProtocol + +from tests.integration import test_scenarios as scenarios + + +@pytest_asyncio.fixture +async def scenario_resources( + monkeypatch: pytest.MonkeyPatch, +) -> AsyncIterator[AsyncExitStack]: + async with AsyncExitStack() as resources: + monkeypatch.setattr(scenarios, '_test_resources', resources) + yield resources + + +async def _assert_server_closed(address: str) -> None: + url = httpx.URL(address if '://' in address else f'http://{address}') + try: + # Avoid combined IPv4/IPv6 refusal errors from localhost resolution. + _, writer = await asyncio.wait_for( + asyncio.open_connection(url.host, url.port, family=socket.AF_INET), + timeout=5, + ) + except ConnectionRefusedError: + return + writer.close() + await writer.wait_closed() + pytest.fail(f'Test server still accepts connections at {address}') + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + 'transport', + [TransportProtocol.GRPC, TransportProtocol.HTTP_JSON], + ids=['grpc', 'rest'], +) +async def test_setup_failure_closes_resources( + monkeypatch: pytest.MonkeyPatch, + scenario_resources: AsyncExitStack, + transport: TransportProtocol, +) -> None: + address = '' + http_client: httpx.AsyncClient | None = None + + def fail_create(factory: ClientFactory, card: AgentCard) -> Client: + nonlocal address, http_client + address = card.supported_interfaces[0].url + http_client = factory._config.httpx_client + raise RuntimeError('client setup failed') + + monkeypatch.setattr(ClientFactory, 'create', fail_create) + with pytest.raises(RuntimeError, match='client setup failed'): + await scenarios.create_client( + scenarios.create_handler(_WaitingAgent(), use_legacy=False), + scenarios.agent_card(), + transport=transport, + ) + + if http_client is not None: + assert http_client.is_closed + await _assert_server_closed(address) + + +@pytest.mark.asyncio +async def test_started_grpc_server_is_closed_when_start_raises( + monkeypatch: pytest.MonkeyPatch, + scenario_resources: AsyncExitStack, +) -> None: + server = grpc.aio.server() + original_start = server.start + original_bind = server.add_insecure_port + address = '' + + def make_server() -> grpc.aio.Server: + return server + + def bind_server(requested_address: str) -> int: + nonlocal address + port = original_bind(requested_address) + address = f'localhost:{port}' + return port + + async def fail_start() -> None: + await original_start() + raise RuntimeError('server startup failed') + + monkeypatch.setattr(scenarios.grpc.aio, 'server', make_server) + monkeypatch.setattr(server, 'add_insecure_port', bind_server) + monkeypatch.setattr(server, 'start', fail_start) + try: + with pytest.raises(RuntimeError, match='server startup failed'): + await scenarios.create_client( + scenarios.create_handler(_WaitingAgent(), use_legacy=False), + scenarios.agent_card(), + ) + await _assert_server_closed(address) + finally: + await server.stop(None) + + +class _WaitingAgent(AgentExecutor): + def __init__(self) -> None: + self.release = asyncio.Event() + self.stopped = asyncio.Event() + + async def execute( + self, context: RequestContext, event_queue: EventQueue + ) -> None: + assert context.message is not None + try: + task = new_task_from_user_message(context.message) + task.status.state = TaskState.TASK_STATE_WORKING + await event_queue.enqueue_event(task) + await self.release.wait() + finally: + self.stopped.set() + + async def cancel( + self, context: RequestContext, event_queue: EventQueue + ) -> None: + self.release.set() + + +@pytest.mark.asyncio +@pytest.mark.timeout(10) +@pytest.mark.parametrize( + 'transport', + [TransportProtocol.GRPC, TransportProtocol.HTTP_JSON], + ids=['grpc', 'rest'], +) +@pytest.mark.parametrize('use_legacy', [False, True], ids=['v2', 'legacy']) +@pytest.mark.parametrize('ending', ['exception', 'cancellation']) +async def test_interrupted_scenario_closes_resources( + scenario_resources: AsyncExitStack, + transport: TransportProtocol, + use_legacy: bool, + ending: str, +) -> None: + agent = _WaitingAgent() + started = asyncio.Event() + address = '' + http_client: httpx.AsyncClient | None = None + stream: AsyncGenerator[StreamResponse, None] | None = None + + async def run_scenario() -> None: + nonlocal address, http_client, stream + async with scenario_resources: + card = scenarios.agent_card() + client = await scenarios.create_client( + scenarios.create_handler(agent, use_legacy), + card, + streaming=True, + transport=transport, + ) + address = card.supported_interfaces[0].url + if transport == TransportProtocol.HTTP_JSON: + http_client = client._transport.httpx_client + stream = client.send_message( + SendMessageRequest( + message=Message( + message_id='cleanup-test', + role=Role.ROLE_USER, + parts=[Part(text='start')], + ) + ) + ) + first = await anext(stream) + assert scenarios.get_state(first) == TaskState.TASK_STATE_WORKING + started.set() + if ending == 'exception': + raise RuntimeError('scenario failed') + await asyncio.Event().wait() + + operation = asyncio.create_task(run_scenario()) + try: + if ending == 'cancellation': + await asyncio.wait_for(started.wait(), timeout=2) + operation.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(operation, timeout=4) + else: + with pytest.raises(RuntimeError, match='scenario failed'): + await asyncio.wait_for(operation, timeout=4) + await asyncio.wait_for(agent.stopped.wait(), timeout=1) + if http_client is not None: + assert http_client.is_closed + await _assert_server_closed(address) + finally: + agent.release.set() + if not operation.done(): + operation.cancel() + await asyncio.gather(operation, return_exceptions=True) + if stream is not None: + await stream.aclose() diff --git a/tests/integration/test_scenarios.py b/tests/integration/test_scenarios.py index 10e86a6a5..1fe98e088 100644 --- a/tests/integration/test_scenarios.py +++ b/tests/integration/test_scenarios.py @@ -1,15 +1,20 @@ import asyncio import collections import logging +import socket +from collections.abc import AsyncIterator +from contextlib import AsyncExitStack, asynccontextmanager from typing import Any import grpc +import httpx import pytest import pytest_asyncio +import uvicorn from a2a.auth.user import User -from a2a.client.client import ClientConfig +from a2a.client.client import Client, ClientConfig from a2a.client.client_factory import ClientFactory from a2a.client.errors import A2AClientError from a2a.helpers.proto_helpers import new_task_from_user_message @@ -21,10 +26,13 @@ DefaultRequestHandlerV2, GrpcHandler, GrpcServerCallContextBuilder, + RequestHandler, ) from a2a.server.request_handlers.default_request_handler import ( LegacyRequestHandler, ) +from a2a.server.routes.common import DefaultServerCallContextBuilder +from a2a.server.routes.rest_routes import create_rest_routes from a2a.server.tasks.inmemory_task_store import InMemoryTaskStore from a2a.types import a2a_pb2_grpc from a2a.types.a2a_pb2 import ( @@ -54,6 +62,8 @@ TaskNotFoundError, UnsupportedOperationError, ) +from starlette.applications import Starlette +from starlette.requests import Request logger = logging.getLogger(__name__) @@ -101,6 +111,12 @@ def build(self, request: Any) -> ServerCallContext: ) +class MockHttpCallContextBuilder(DefaultServerCallContextBuilder): + def build_user(self, request: Request) -> User: + """Uses the same test identity as the gRPC setup.""" + return MockUser() + + class InputRequiredAgent(AgentExecutor): """Test agent: 'start' -> INPUT_REQUIRED, anything else -> COMPLETED.""" @@ -167,43 +183,128 @@ def validate_state(event, expected_state): assert get_state(event) == expected_state -_test_servers = [] +_test_resources = AsyncExitStack() @pytest_asyncio.fixture(autouse=True) -async def cleanup_test_servers(): - yield - for server in _test_servers: - await server.stop(None) - _test_servers.clear() - - -# TODO: Test different transport (e.g. HTTP_JSON hangs for some tests). -async def create_client(handler, agent_card, streaming=False): - server = grpc.aio.server() - port = server.add_insecure_port('[::]:0') - server_address = f'localhost:{port}' - - agent_card.supported_interfaces[0].url = server_address - agent_card.supported_interfaces[0].protocol_binding = TransportProtocol.GRPC - - servicer = GrpcHandler( - request_handler=handler, context_builder=MockCallContextBuilder() - ) - a2a_pb2_grpc.add_A2AServiceServicer_to_server(servicer, server) - await server.start() - _test_servers.append(server) - - factory = ClientFactory( - config=ClientConfig( - grpc_channel_factory=grpc.aio.insecure_channel, - supported_protocol_bindings=[TransportProtocol.GRPC], - streaming=streaming, +async def cleanup_test_servers() -> AsyncIterator[None]: + try: + yield + finally: + await _test_resources.aclose() + + +async def close_scenario_handler(handler: RequestHandler) -> None: + """Stops agent work owned by a scenario, including legacy background work.""" + if isinstance(handler, DefaultRequestHandlerV2): + await handler.aclose() + elif isinstance(handler, LegacyRequestHandler): + with handler._running_agents_lock: + producers = list(handler._running_agents.values()) + for producer in producers: + producer.cancel() + if producers: + await asyncio.gather(*producers, return_exceptions=True) + background_tasks = list(handler._background_tasks) + if background_tasks: + await asyncio.gather(*background_tasks, return_exceptions=True) + + +@asynccontextmanager +async def running_rest_server(app: Starlette) -> AsyncIterator[str]: + """Runs a local HTTP server for streaming scenario tests.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as listener: + listener.bind(('127.0.0.1', 0)) + port = listener.getsockname()[1] + server = uvicorn.Server( + uvicorn.Config( + app, + host='127.0.0.1', + port=port, + lifespan='off', + log_level='warning', + timeout_graceful_shutdown=1, + ) ) - ) - client = factory.create(agent_card) - client._server = server # Keep reference to prevent garbage collection - return client + server_task = asyncio.create_task(server.serve(sockets=[listener])) + try: + deadline = asyncio.get_running_loop().time() + 5 + # Uvicorn exposes readiness as a flag rather than an event. + while not server.started: # noqa: ASYNC110 + if server_task.done(): + await server_task + raise RuntimeError('REST test server exited before startup') + if asyncio.get_running_loop().time() >= deadline: + raise TimeoutError('REST test server did not start') + await asyncio.sleep(0.01) + yield f'http://127.0.0.1:{port}' + finally: + server.should_exit = True + try: + await asyncio.wait_for(server_task, timeout=5) + except asyncio.TimeoutError: + server_task.cancel() + await asyncio.gather(server_task, return_exceptions=True) + + +async def create_client( + handler: RequestHandler, + agent_card: AgentCard, + streaming: bool = False, + *, + transport: TransportProtocol = TransportProtocol.GRPC, +) -> Client: + """Creates a client for a scenario test.""" + if transport not in (TransportProtocol.GRPC, TransportProtocol.HTTP_JSON): + raise ValueError(f'Unsupported scenario transport: {transport}') + + async with AsyncExitStack() as resources: + resources.push_async_callback(close_scenario_handler, handler) + + if transport == TransportProtocol.HTTP_JSON: + app = Starlette( + routes=create_rest_routes( + request_handler=handler, + context_builder=MockHttpCallContextBuilder(), + ) + ) + server_address = await resources.enter_async_context( + running_rest_server(app) + ) + httpx_client = await resources.enter_async_context( + httpx.AsyncClient(trust_env=False) + ) + config = ClientConfig( + httpx_client=httpx_client, + supported_protocol_bindings=[transport], + streaming=streaming, + ) + else: + server = grpc.aio.server() + resources.push_async_callback(server.stop, None) + port = server.add_insecure_port('[::]:0') + server_address = f'localhost:{port}' + + servicer = GrpcHandler( + request_handler=handler, + context_builder=MockCallContextBuilder(), + ) + a2a_pb2_grpc.add_A2AServiceServicer_to_server(servicer, server) + await server.start() + config = ClientConfig( + grpc_channel_factory=grpc.aio.insecure_channel, + supported_protocol_bindings=[transport], + streaming=streaming, + ) + + agent_card.supported_interfaces[0].url = server_address + agent_card.supported_interfaces[0].protocol_binding = transport + + factory = ClientFactory(config=config) + client = factory.create(agent_card) + resources.push_async_callback(client.close) + _test_resources.push_async_callback(resources.pop_all().aclose) + return client def create_handler( @@ -232,20 +333,27 @@ def create_handler( # This also covers test_scenario_7_cancel_terminal_task from test_handler_comparison @pytest.mark.timeout(2.0) @pytest.mark.asyncio +@pytest.mark.parametrize( + 'transport', + [TransportProtocol.GRPC, TransportProtocol.HTTP_JSON], + ids=['grpc', 'rest'], +) @pytest.mark.parametrize('use_legacy', [False, True], ids=['v2', 'legacy']) @pytest.mark.parametrize( 'streaming', [False, True], ids=['blocking', 'streaming'] ) -async def test_scenario_1_cancel_terminal_task(use_legacy, streaming): +async def test_scenario_1_cancel_terminal_task( + use_legacy: bool, streaming: bool, transport: TransportProtocol +) -> None: class DummyAgentExecutor(AgentExecutor): async def execute( self, context: RequestContext, event_queue: EventQueue - ): + ) -> None: pass async def cancel( self, context: RequestContext, event_queue: EventQueue - ): + ) -> None: pass task_store = InMemoryTaskStore() @@ -253,7 +361,10 @@ async def cancel( DummyAgentExecutor(), use_legacy, task_store=task_store ) client = await create_client( - handler, agent_card=agent_card(), streaming=streaming + handler, + agent_card=agent_card(), + streaming=streaming, + transport=transport, ) task_id = 'terminal-task' @@ -268,12 +379,19 @@ async def cancel( @pytest.mark.asyncio +@pytest.mark.parametrize( + 'transport', + [TransportProtocol.GRPC, TransportProtocol.HTTP_JSON], + ids=['grpc', 'rest'], +) @pytest.mark.parametrize('use_legacy', [False, True], ids=['v2', 'legacy']) -async def test_scenario_4_simple_streaming(use_legacy): +async def test_scenario_4_simple_streaming( + use_legacy: bool, transport: TransportProtocol +) -> None: class DummyAgentExecutor(AgentExecutor): async def execute( self, context: RequestContext, event_queue: EventQueue - ): + ) -> None: task = new_task_from_user_message(context.message) task.status.state = TaskState.TASK_STATE_WORKING await event_queue.enqueue_event(task) @@ -287,12 +405,12 @@ async def execute( async def cancel( self, context: RequestContext, event_queue: EventQueue - ): + ) -> None: pass handler = create_handler(DummyAgentExecutor(), use_legacy) client = await create_client( - handler, agent_card=agent_card(), streaming=True + handler, agent_card=agent_card(), streaming=True, transport=transport ) msg = Message( message_id='test-msg', role=Role.ROLE_USER, parts=[Part(text='hello')] @@ -1576,15 +1694,22 @@ async def cancel( # Scenario: Publish ArtifactUpdateEvent @pytest.mark.timeout(2.0) @pytest.mark.asyncio +@pytest.mark.parametrize( + 'transport', + [TransportProtocol.GRPC, TransportProtocol.HTTP_JSON], + ids=['grpc', 'rest'], +) @pytest.mark.parametrize('use_legacy', [False, True], ids=['v2', 'legacy']) @pytest.mark.parametrize( 'streaming', [False, True], ids=['blocking', 'streaming'] ) -async def test_scenario_publish_artifact(use_legacy, streaming): +async def test_scenario_publish_artifact( + use_legacy: bool, streaming: bool, transport: TransportProtocol +) -> None: class ArtifactAgent(AgentExecutor): async def execute( self, context: RequestContext, event_queue: EventQueue - ): + ) -> None: task = new_task_from_user_message(context.message) task.status.state = TaskState.TASK_STATE_WORKING await event_queue.enqueue_event(task) @@ -1607,12 +1732,15 @@ async def execute( async def cancel( self, context: RequestContext, event_queue: EventQueue - ): + ) -> None: pass handler = create_handler(ArtifactAgent(), use_legacy) client = await create_client( - handler, agent_card=agent_card(), streaming=streaming + handler, + agent_card=agent_card(), + streaming=streaming, + transport=transport, ) msg = Message( From 6a6df820fbccbee4bd2766c174ae297fbb8b9c29 Mon Sep 17 00:00:00 2001 From: Faheem Shah Umer Vattam Kandathil <161448038+faheem-shah-umer@users.noreply.github.com> Date: Mon, 26 Oct 2026 14:13:49 +0100 Subject: [PATCH 2/2] test: simplify REST scenario setup after review --- tests/integration/test_scenario_cleanup.py | 218 --------------------- tests/integration/test_scenarios.py | 98 ++++----- 2 files changed, 42 insertions(+), 274 deletions(-) delete mode 100644 tests/integration/test_scenario_cleanup.py diff --git a/tests/integration/test_scenario_cleanup.py b/tests/integration/test_scenario_cleanup.py deleted file mode 100644 index 66a8b51f3..000000000 --- a/tests/integration/test_scenario_cleanup.py +++ /dev/null @@ -1,218 +0,0 @@ -import asyncio -import socket - -from collections.abc import AsyncGenerator, AsyncIterator -from contextlib import AsyncExitStack - -import grpc -import httpx -import pytest -import pytest_asyncio - -from a2a.client.client import Client -from a2a.client.client_factory import ClientFactory -from a2a.helpers.proto_helpers import new_task_from_user_message -from a2a.server.agent_execution import AgentExecutor, RequestContext -from a2a.server.events import EventQueue -from a2a.types.a2a_pb2 import ( - AgentCard, - Message, - Part, - Role, - SendMessageRequest, - StreamResponse, - TaskState, -) -from a2a.utils import TransportProtocol - -from tests.integration import test_scenarios as scenarios - - -@pytest_asyncio.fixture -async def scenario_resources( - monkeypatch: pytest.MonkeyPatch, -) -> AsyncIterator[AsyncExitStack]: - async with AsyncExitStack() as resources: - monkeypatch.setattr(scenarios, '_test_resources', resources) - yield resources - - -async def _assert_server_closed(address: str) -> None: - url = httpx.URL(address if '://' in address else f'http://{address}') - try: - # Avoid combined IPv4/IPv6 refusal errors from localhost resolution. - _, writer = await asyncio.wait_for( - asyncio.open_connection(url.host, url.port, family=socket.AF_INET), - timeout=5, - ) - except ConnectionRefusedError: - return - writer.close() - await writer.wait_closed() - pytest.fail(f'Test server still accepts connections at {address}') - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - 'transport', - [TransportProtocol.GRPC, TransportProtocol.HTTP_JSON], - ids=['grpc', 'rest'], -) -async def test_setup_failure_closes_resources( - monkeypatch: pytest.MonkeyPatch, - scenario_resources: AsyncExitStack, - transport: TransportProtocol, -) -> None: - address = '' - http_client: httpx.AsyncClient | None = None - - def fail_create(factory: ClientFactory, card: AgentCard) -> Client: - nonlocal address, http_client - address = card.supported_interfaces[0].url - http_client = factory._config.httpx_client - raise RuntimeError('client setup failed') - - monkeypatch.setattr(ClientFactory, 'create', fail_create) - with pytest.raises(RuntimeError, match='client setup failed'): - await scenarios.create_client( - scenarios.create_handler(_WaitingAgent(), use_legacy=False), - scenarios.agent_card(), - transport=transport, - ) - - if http_client is not None: - assert http_client.is_closed - await _assert_server_closed(address) - - -@pytest.mark.asyncio -async def test_started_grpc_server_is_closed_when_start_raises( - monkeypatch: pytest.MonkeyPatch, - scenario_resources: AsyncExitStack, -) -> None: - server = grpc.aio.server() - original_start = server.start - original_bind = server.add_insecure_port - address = '' - - def make_server() -> grpc.aio.Server: - return server - - def bind_server(requested_address: str) -> int: - nonlocal address - port = original_bind(requested_address) - address = f'localhost:{port}' - return port - - async def fail_start() -> None: - await original_start() - raise RuntimeError('server startup failed') - - monkeypatch.setattr(scenarios.grpc.aio, 'server', make_server) - monkeypatch.setattr(server, 'add_insecure_port', bind_server) - monkeypatch.setattr(server, 'start', fail_start) - try: - with pytest.raises(RuntimeError, match='server startup failed'): - await scenarios.create_client( - scenarios.create_handler(_WaitingAgent(), use_legacy=False), - scenarios.agent_card(), - ) - await _assert_server_closed(address) - finally: - await server.stop(None) - - -class _WaitingAgent(AgentExecutor): - def __init__(self) -> None: - self.release = asyncio.Event() - self.stopped = asyncio.Event() - - async def execute( - self, context: RequestContext, event_queue: EventQueue - ) -> None: - assert context.message is not None - try: - task = new_task_from_user_message(context.message) - task.status.state = TaskState.TASK_STATE_WORKING - await event_queue.enqueue_event(task) - await self.release.wait() - finally: - self.stopped.set() - - async def cancel( - self, context: RequestContext, event_queue: EventQueue - ) -> None: - self.release.set() - - -@pytest.mark.asyncio -@pytest.mark.timeout(10) -@pytest.mark.parametrize( - 'transport', - [TransportProtocol.GRPC, TransportProtocol.HTTP_JSON], - ids=['grpc', 'rest'], -) -@pytest.mark.parametrize('use_legacy', [False, True], ids=['v2', 'legacy']) -@pytest.mark.parametrize('ending', ['exception', 'cancellation']) -async def test_interrupted_scenario_closes_resources( - scenario_resources: AsyncExitStack, - transport: TransportProtocol, - use_legacy: bool, - ending: str, -) -> None: - agent = _WaitingAgent() - started = asyncio.Event() - address = '' - http_client: httpx.AsyncClient | None = None - stream: AsyncGenerator[StreamResponse, None] | None = None - - async def run_scenario() -> None: - nonlocal address, http_client, stream - async with scenario_resources: - card = scenarios.agent_card() - client = await scenarios.create_client( - scenarios.create_handler(agent, use_legacy), - card, - streaming=True, - transport=transport, - ) - address = card.supported_interfaces[0].url - if transport == TransportProtocol.HTTP_JSON: - http_client = client._transport.httpx_client - stream = client.send_message( - SendMessageRequest( - message=Message( - message_id='cleanup-test', - role=Role.ROLE_USER, - parts=[Part(text='start')], - ) - ) - ) - first = await anext(stream) - assert scenarios.get_state(first) == TaskState.TASK_STATE_WORKING - started.set() - if ending == 'exception': - raise RuntimeError('scenario failed') - await asyncio.Event().wait() - - operation = asyncio.create_task(run_scenario()) - try: - if ending == 'cancellation': - await asyncio.wait_for(started.wait(), timeout=2) - operation.cancel() - with pytest.raises(asyncio.CancelledError): - await asyncio.wait_for(operation, timeout=4) - else: - with pytest.raises(RuntimeError, match='scenario failed'): - await asyncio.wait_for(operation, timeout=4) - await asyncio.wait_for(agent.stopped.wait(), timeout=1) - if http_client is not None: - assert http_client.is_closed - await _assert_server_closed(address) - finally: - agent.release.set() - if not operation.done(): - operation.cancel() - await asyncio.gather(operation, return_exceptions=True) - if stream is not None: - await stream.aclose() diff --git a/tests/integration/test_scenarios.py b/tests/integration/test_scenarios.py index 1fe98e088..797e0cf8d 100644 --- a/tests/integration/test_scenarios.py +++ b/tests/integration/test_scenarios.py @@ -1,17 +1,15 @@ import asyncio import collections import logging -import socket from collections.abc import AsyncIterator -from contextlib import AsyncExitStack, asynccontextmanager -from typing import Any +from contextlib import AsyncExitStack +from typing import Any, NamedTuple import grpc import httpx import pytest import pytest_asyncio -import uvicorn from a2a.auth.user import User from a2a.client.client import Client, ClientConfig @@ -183,6 +181,20 @@ def validate_state(event, expected_state): assert get_state(event) == expected_state +class RestSetup(NamedTuple): + app: Starlette + httpx_client: httpx.AsyncClient + + +@pytest_asyncio.fixture(scope='module') +async def rest_setup() -> AsyncIterator[RestSetup]: + app = Starlette() + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), base_url='http://testserver' + ) as httpx_client: + yield RestSetup(app, httpx_client) + + _test_resources = AsyncExitStack() @@ -210,49 +222,13 @@ async def close_scenario_handler(handler: RequestHandler) -> None: await asyncio.gather(*background_tasks, return_exceptions=True) -@asynccontextmanager -async def running_rest_server(app: Starlette) -> AsyncIterator[str]: - """Runs a local HTTP server for streaming scenario tests.""" - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as listener: - listener.bind(('127.0.0.1', 0)) - port = listener.getsockname()[1] - server = uvicorn.Server( - uvicorn.Config( - app, - host='127.0.0.1', - port=port, - lifespan='off', - log_level='warning', - timeout_graceful_shutdown=1, - ) - ) - server_task = asyncio.create_task(server.serve(sockets=[listener])) - try: - deadline = asyncio.get_running_loop().time() + 5 - # Uvicorn exposes readiness as a flag rather than an event. - while not server.started: # noqa: ASYNC110 - if server_task.done(): - await server_task - raise RuntimeError('REST test server exited before startup') - if asyncio.get_running_loop().time() >= deadline: - raise TimeoutError('REST test server did not start') - await asyncio.sleep(0.01) - yield f'http://127.0.0.1:{port}' - finally: - server.should_exit = True - try: - await asyncio.wait_for(server_task, timeout=5) - except asyncio.TimeoutError: - server_task.cancel() - await asyncio.gather(server_task, return_exceptions=True) - - async def create_client( handler: RequestHandler, agent_card: AgentCard, streaming: bool = False, *, transport: TransportProtocol = TransportProtocol.GRPC, + rest_setup: RestSetup | None = None, ) -> Client: """Creates a client for a scenario test.""" if transport not in (TransportProtocol.GRPC, TransportProtocol.HTTP_JSON): @@ -262,20 +238,18 @@ async def create_client( resources.push_async_callback(close_scenario_handler, handler) if transport == TransportProtocol.HTTP_JSON: - app = Starlette( - routes=create_rest_routes( - request_handler=handler, - context_builder=MockHttpCallContextBuilder(), + if rest_setup is None: + raise ValueError( + 'REST scenarios require the rest_setup fixture' ) + rest_setup.app.router.routes[:] = create_rest_routes( + request_handler=handler, + context_builder=MockHttpCallContextBuilder(), ) - server_address = await resources.enter_async_context( - running_rest_server(app) - ) - httpx_client = await resources.enter_async_context( - httpx.AsyncClient(trust_env=False) - ) + resources.callback(rest_setup.app.router.routes.clear) + server_address = 'http://testserver' config = ClientConfig( - httpx_client=httpx_client, + httpx_client=rest_setup.httpx_client, supported_protocol_bindings=[transport], streaming=streaming, ) @@ -343,7 +317,10 @@ def create_handler( 'streaming', [False, True], ids=['blocking', 'streaming'] ) async def test_scenario_1_cancel_terminal_task( - use_legacy: bool, streaming: bool, transport: TransportProtocol + use_legacy: bool, + streaming: bool, + transport: TransportProtocol, + rest_setup: RestSetup, ) -> None: class DummyAgentExecutor(AgentExecutor): async def execute( @@ -365,6 +342,7 @@ async def cancel( agent_card=agent_card(), streaming=streaming, transport=transport, + rest_setup=rest_setup, ) task_id = 'terminal-task' @@ -386,7 +364,7 @@ async def cancel( ) @pytest.mark.parametrize('use_legacy', [False, True], ids=['v2', 'legacy']) async def test_scenario_4_simple_streaming( - use_legacy: bool, transport: TransportProtocol + use_legacy: bool, transport: TransportProtocol, rest_setup: RestSetup ) -> None: class DummyAgentExecutor(AgentExecutor): async def execute( @@ -410,7 +388,11 @@ async def cancel( handler = create_handler(DummyAgentExecutor(), use_legacy) client = await create_client( - handler, agent_card=agent_card(), streaming=True, transport=transport + handler, + agent_card=agent_card(), + streaming=True, + transport=transport, + rest_setup=rest_setup, ) msg = Message( message_id='test-msg', role=Role.ROLE_USER, parts=[Part(text='hello')] @@ -1704,7 +1686,10 @@ async def cancel( 'streaming', [False, True], ids=['blocking', 'streaming'] ) async def test_scenario_publish_artifact( - use_legacy: bool, streaming: bool, transport: TransportProtocol + use_legacy: bool, + streaming: bool, + transport: TransportProtocol, + rest_setup: RestSetup, ) -> None: class ArtifactAgent(AgentExecutor): async def execute( @@ -1741,6 +1726,7 @@ async def cancel( agent_card=agent_card(), streaming=streaming, transport=transport, + rest_setup=rest_setup, ) msg = Message(