mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
fix: Retain the consumers and transport registry when recreating the ClientFactory in remote_a2a_agent.py
PiperOrigin-RevId: 807762203
This commit is contained in:
committed by
Copybara-Service
parent
f7bd3c111c
commit
6bd33e1be3
@@ -188,12 +188,16 @@ class RemoteA2aAgent(BaseAgent):
|
||||
)
|
||||
self._httpx_client_needs_cleanup = True
|
||||
if self._a2a_client_factory:
|
||||
registry = self._a2a_client_factory._registry
|
||||
self._a2a_client_factory = A2AClientFactory(
|
||||
config=dataclasses.replace(
|
||||
self._a2a_client_factory._config,
|
||||
httpx_client=self._httpx_client,
|
||||
)
|
||||
),
|
||||
consumers=self._a2a_client_factory._consumers,
|
||||
)
|
||||
for label, generator in registry.items():
|
||||
self._a2a_client_factory.register(label, generator)
|
||||
if not self._a2a_client_factory:
|
||||
client_config = A2AClientConfig(
|
||||
httpx_client=self._httpx_client,
|
||||
|
||||
@@ -34,6 +34,7 @@ pytestmark = pytest.mark.skipif(
|
||||
# Import dependencies with version checking
|
||||
try:
|
||||
from a2a.client.client import ClientConfig
|
||||
from a2a.client.client import Consumer
|
||||
from a2a.client.client_factory import ClientFactory
|
||||
from a2a.types import AgentCapabilities
|
||||
from a2a.types import AgentCard
|
||||
@@ -246,6 +247,50 @@ class TestRemoteA2aAgentResolution:
|
||||
assert client == existing_client
|
||||
assert agent._httpx_client_needs_cleanup is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ensure_httpx_client_updates_factory_with_new_client(self):
|
||||
"""Test that _ensure_httpx_client updates factory with new client."""
|
||||
agent = RemoteA2aAgent(
|
||||
name="test_agent",
|
||||
agent_card=create_test_agent_card(),
|
||||
a2a_client_factory=ClientFactory(
|
||||
ClientConfig(httpx_client=None),
|
||||
),
|
||||
)
|
||||
assert agent._a2a_client_factory._config.httpx_client is None
|
||||
|
||||
client = await agent._ensure_httpx_client()
|
||||
|
||||
assert client is not None
|
||||
assert agent._httpx_client == client
|
||||
assert agent._httpx_client_needs_cleanup is True
|
||||
assert agent._a2a_client_factory._config.httpx_client == client
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ensure_httpx_client_reregisters_transports_with_new_client(
|
||||
self,
|
||||
):
|
||||
"""Test that _ensure_httpx_client registers transports with new client."""
|
||||
factory = ClientFactory(
|
||||
ClientConfig(httpx_client=None),
|
||||
)
|
||||
factory.register("transport_label", lambda: "test")
|
||||
agent = RemoteA2aAgent(
|
||||
name="test_agent",
|
||||
agent_card=create_test_agent_card(),
|
||||
a2a_client_factory=factory,
|
||||
)
|
||||
assert agent._a2a_client_factory._config.httpx_client is None
|
||||
assert "transport_label" in agent._a2a_client_factory._registry
|
||||
|
||||
client = await agent._ensure_httpx_client()
|
||||
|
||||
assert client is not None
|
||||
assert agent._httpx_client == client
|
||||
assert agent._httpx_client_needs_cleanup is True
|
||||
assert agent._a2a_client_factory._config.httpx_client == client
|
||||
assert "transport_label" in agent._a2a_client_factory._registry
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_agent_card_from_url_success(self):
|
||||
"""Test successful agent card resolution from URL."""
|
||||
|
||||
Reference in New Issue
Block a user