Skip to content

Commit 08ef9d8

Browse files
committed
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.
1 parent ca7edc3 commit 08ef9d8

3 files changed

Lines changed: 22 additions & 1 deletion

File tree

src/a2a/client/transports/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
from a2a.client.transports.base import ClientTransport
44
from a2a.client.transports.jsonrpc import JsonRpcTransport
55
from a2a.client.transports.rest import RestTransport
6+
from a2a.client.transports.tenant_decorator import TenantTransportDecorator
67

78

89
try:
@@ -16,4 +17,5 @@
1617
'GrpcTransport',
1718
'JsonRpcTransport',
1819
'RestTransport',
20+
'TenantTransportDecorator',
1921
]

src/a2a/client/transports/tenant_decorator.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ async def send_message(
4343
*,
4444
context: ClientCallContext | None = None,
4545
) -> SendMessageResponse:
46-
"""Sends a streaming message request to the agent and yields responses as they arrive."""
46+
"""Sends a non-streaming message request to the agent."""
4747
request.tenant = self._resolve_tenant(request.tenant)
4848
return await self._base.send_message(request, context=context)
4949

tests/client/transports/test_tenant_decorator.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -127,3 +127,22 @@ async def mock_stream(*args, **kwargs):
127127
async for _ in decorator.send_message_streaming(request_msg):
128128
pass
129129
assert request_msg.tenant == tenant_id
130+
131+
@pytest.mark.asyncio
132+
async def test_close_delegates_to_base(
133+
self, mock_transport: AsyncMock
134+
) -> None:
135+
"""Test that close() is delegated to the underlying transport."""
136+
decorator = TenantTransportDecorator(mock_transport, 'test-tenant')
137+
await decorator.close()
138+
mock_transport.close.assert_awaited_once()
139+
140+
@pytest.mark.asyncio
141+
async def test_async_context_manager(
142+
self, mock_transport: AsyncMock
143+
) -> None:
144+
"""Test that the decorator works as an async context manager."""
145+
decorator = TenantTransportDecorator(mock_transport, 'test-tenant')
146+
async with decorator as transport:
147+
assert transport is decorator
148+
mock_transport.close.assert_awaited_once()

0 commit comments

Comments
 (0)