mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
feat: Allow injecting a custom Runner into agent_to_a2a
This change adds an optional `runner` parameter to the `to_a2a` function, enabling users to provide a pre-configured `Runner` instance instead of always using the default in-memory services. A new test case has been added to verify this functionality. closes #3104 Co-authored-by: Dongyu Jia <dongyuj@google.com> PiperOrigin-RevId: 826526861
This commit is contained in:
committed by
Copybara-Service
parent
338c3c89c6
commit
156d235479
@@ -103,6 +103,51 @@ class TestToA2A:
|
||||
"startup", mock_app.add_event_handler.call_args[0][1]
|
||||
)
|
||||
|
||||
@patch("google.adk.a2a.utils.agent_to_a2a.A2aAgentExecutor")
|
||||
@patch("google.adk.a2a.utils.agent_to_a2a.DefaultRequestHandler")
|
||||
@patch("google.adk.a2a.utils.agent_to_a2a.InMemoryTaskStore")
|
||||
@patch("google.adk.a2a.utils.agent_to_a2a.AgentCardBuilder")
|
||||
@patch("google.adk.a2a.utils.agent_to_a2a.Starlette")
|
||||
def test_to_a2a_with_custom_runner(
|
||||
self,
|
||||
mock_starlette_class,
|
||||
mock_card_builder_class,
|
||||
mock_task_store_class,
|
||||
mock_request_handler_class,
|
||||
mock_agent_executor_class,
|
||||
):
|
||||
"""Test to_a2a with a custom runner."""
|
||||
# Arrange
|
||||
mock_app = Mock(spec=Starlette)
|
||||
mock_starlette_class.return_value = mock_app
|
||||
mock_task_store = Mock(spec=InMemoryTaskStore)
|
||||
mock_task_store_class.return_value = mock_task_store
|
||||
mock_agent_executor = Mock(spec=A2aAgentExecutor)
|
||||
mock_agent_executor_class.return_value = mock_agent_executor
|
||||
mock_request_handler = Mock(spec=DefaultRequestHandler)
|
||||
mock_request_handler_class.return_value = mock_request_handler
|
||||
mock_card_builder = Mock(spec=AgentCardBuilder)
|
||||
mock_card_builder_class.return_value = mock_card_builder
|
||||
custom_runner = Mock(spec=Runner)
|
||||
|
||||
# Act
|
||||
result = to_a2a(self.mock_agent, runner=custom_runner)
|
||||
|
||||
# Assert
|
||||
assert result == mock_app
|
||||
mock_starlette_class.assert_called_once()
|
||||
mock_task_store_class.assert_called_once()
|
||||
mock_agent_executor_class.assert_called_once_with(runner=custom_runner)
|
||||
mock_request_handler_class.assert_called_once_with(
|
||||
agent_executor=mock_agent_executor, task_store=mock_task_store
|
||||
)
|
||||
mock_card_builder_class.assert_called_once_with(
|
||||
agent=self.mock_agent, rpc_url="http://localhost:8000/"
|
||||
)
|
||||
mock_app.add_event_handler.assert_called_once_with(
|
||||
"startup", mock_app.add_event_handler.call_args[0][1]
|
||||
)
|
||||
|
||||
@patch("google.adk.a2a.utils.agent_to_a2a.A2aAgentExecutor")
|
||||
@patch("google.adk.a2a.utils.agent_to_a2a.DefaultRequestHandler")
|
||||
@patch("google.adk.a2a.utils.agent_to_a2a.InMemoryTaskStore")
|
||||
|
||||
Reference in New Issue
Block a user