From 08ef9d88287a19a14bf79763ee86ac33bdbf3ff3 Mon Sep 17 00:00:00 2001 From: 4444jPPP Date: Mon, 30 Mar 2026 17:45:28 -0400 Subject: [PATCH] fix(client): export TenantTransportDecorator and fix docstring - Add TenantTransportDecorator to transports __init__.py __all__ so it is importable via `from a2a.client.transports import TenantTransportDecorator` alongside the other transports. - Fix send_message() docstring that incorrectly described it as a streaming method (copy-paste from send_message_streaming). - Add test coverage for close() delegation and async context manager usage, verifying the decorator properly delegates lifecycle management to the underlying transport. --- src/a2a/client/transports/__init__.py | 2 ++ src/a2a/client/transports/tenant_decorator.py | 2 +- .../transports/test_tenant_decorator.py | 19 +++++++++++++++++++ 3 files changed, 22 insertions(+), 1 deletion(-) diff --git a/src/a2a/client/transports/__init__.py b/src/a2a/client/transports/__init__.py index af7c60f62..001ca67b5 100644 --- a/src/a2a/client/transports/__init__.py +++ b/src/a2a/client/transports/__init__.py @@ -3,6 +3,7 @@ from a2a.client.transports.base import ClientTransport from a2a.client.transports.jsonrpc import JsonRpcTransport from a2a.client.transports.rest import RestTransport +from a2a.client.transports.tenant_decorator import TenantTransportDecorator try: @@ -16,4 +17,5 @@ 'GrpcTransport', 'JsonRpcTransport', 'RestTransport', + 'TenantTransportDecorator', ] diff --git a/src/a2a/client/transports/tenant_decorator.py b/src/a2a/client/transports/tenant_decorator.py index d1059d757..80f596d2e 100644 --- a/src/a2a/client/transports/tenant_decorator.py +++ b/src/a2a/client/transports/tenant_decorator.py @@ -43,7 +43,7 @@ async def send_message( *, context: ClientCallContext | None = None, ) -> SendMessageResponse: - """Sends a streaming message request to the agent and yields responses as they arrive.""" + """Sends a non-streaming message request to the agent.""" request.tenant = self._resolve_tenant(request.tenant) return await self._base.send_message(request, context=context) diff --git a/tests/client/transports/test_tenant_decorator.py b/tests/client/transports/test_tenant_decorator.py index b08406bad..1e560d2ac 100644 --- a/tests/client/transports/test_tenant_decorator.py +++ b/tests/client/transports/test_tenant_decorator.py @@ -127,3 +127,22 @@ async def mock_stream(*args, **kwargs): async for _ in decorator.send_message_streaming(request_msg): pass assert request_msg.tenant == tenant_id + + @pytest.mark.asyncio + async def test_close_delegates_to_base( + self, mock_transport: AsyncMock + ) -> None: + """Test that close() is delegated to the underlying transport.""" + decorator = TenantTransportDecorator(mock_transport, 'test-tenant') + await decorator.close() + mock_transport.close.assert_awaited_once() + + @pytest.mark.asyncio + async def test_async_context_manager( + self, mock_transport: AsyncMock + ) -> None: + """Test that the decorator works as an async context manager.""" + decorator = TenantTransportDecorator(mock_transport, 'test-tenant') + async with decorator as transport: + assert transport is decorator + mock_transport.close.assert_awaited_once()