Move public_utils to utils in tests

Renamed conflicting utils.py as testing_utils.py

PiperOrigin-RevId: 761715808
This commit is contained in:
Selcuk Gun 2025-05-21 16:34:28 -07:00 committed by Copybara-Service
parent 09cb128cf9
commit 41b33d4a0a
22 changed files with 236 additions and 208 deletions

View File

@ -14,7 +14,7 @@
from pytest import mark from pytest import mark
from ..unittests.utils import simplify_events from ..unittests.testing_utils import simplify_events
from .fixture import callback_agent from .fixture import callback_agent
from .utils import assert_agent_says from .utils import assert_agent_says
from .utils import TestRunner from .utils import TestRunner

View File

@ -30,7 +30,7 @@ from google.genai import types
import pytest import pytest
import pytest_mock import pytest_mock
from typing_extensions import override from typing_extensions import override
from .. import utils from .. import testing_utils
def _before_agent_callback_noop(callback_context: CallbackContext) -> None: def _before_agent_callback_noop(callback_context: CallbackContext) -> None:
@ -398,7 +398,7 @@ async def test_before_agent_callbacks_chain(
request.function.__name__, agent request.function.__name__, agent
) )
result = [e async for e in agent.run_async(parent_ctx)] result = [e async for e in agent.run_async(parent_ctx)]
assert utils.simplify_events(result) == [ assert testing_utils.simplify_events(result) == [
(f'{request.function.__name__}_test_agent', response) (f'{request.function.__name__}_test_agent', response)
for response in expected_responses for response in expected_responses
] ]
@ -459,7 +459,7 @@ async def test_after_agent_callbacks_chain(
request.function.__name__, agent request.function.__name__, agent
) )
result = [e async for e in agent.run_async(parent_ctx)] result = [e async for e in agent.run_async(parent_ctx)]
assert utils.simplify_events(result) == [ assert testing_utils.simplify_events(result) == [
(f'{request.function.__name__}_test_agent', response) (f'{request.function.__name__}_test_agent', response)
for response in expected_responses for response in expected_responses
] ]

View File

@ -23,7 +23,7 @@ from google.genai import types
from pydantic import BaseModel from pydantic import BaseModel
import pytest import pytest
from .. import utils from .. import testing_utils
class MockBeforeModelCallback(BaseModel): class MockBeforeModelCallback(BaseModel):
@ -35,7 +35,7 @@ class MockBeforeModelCallback(BaseModel):
llm_request: LlmRequest, llm_request: LlmRequest,
) -> LlmResponse: ) -> LlmResponse:
return LlmResponse( return LlmResponse(
content=utils.ModelContent( content=testing_utils.ModelContent(
[types.Part.from_text(text=self.mock_response)] [types.Part.from_text(text=self.mock_response)]
) )
) )
@ -50,7 +50,7 @@ class MockAfterModelCallback(BaseModel):
llm_response: LlmResponse, llm_response: LlmResponse,
) -> LlmResponse: ) -> LlmResponse:
return LlmResponse( return LlmResponse(
content=utils.ModelContent( content=testing_utils.ModelContent(
[types.Part.from_text(text=self.mock_response)] [types.Part.from_text(text=self.mock_response)]
) )
) )
@ -65,7 +65,7 @@ class MockAsyncBeforeModelCallback(BaseModel):
llm_request: LlmRequest, llm_request: LlmRequest,
) -> LlmResponse: ) -> LlmResponse:
return LlmResponse( return LlmResponse(
content=utils.ModelContent( content=testing_utils.ModelContent(
[types.Part.from_text(text=self.mock_response)] [types.Part.from_text(text=self.mock_response)]
) )
) )
@ -80,7 +80,7 @@ class MockAsyncAfterModelCallback(BaseModel):
llm_response: LlmResponse, llm_response: LlmResponse,
) -> LlmResponse: ) -> LlmResponse:
return LlmResponse( return LlmResponse(
content=utils.ModelContent( content=testing_utils.ModelContent(
[types.Part.from_text(text=self.mock_response)] [types.Part.from_text(text=self.mock_response)]
) )
) )
@ -97,7 +97,7 @@ async def async_noop_callback(**kwargs) -> Optional[LlmResponse]:
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_before_model_callback(): async def test_before_model_callback():
responses = ['model_response'] responses = ['model_response']
mock_model = utils.MockModel.create(responses=responses) mock_model = testing_utils.MockModel.create(responses=responses)
agent = Agent( agent = Agent(
name='root_agent', name='root_agent',
model=mock_model, model=mock_model,
@ -106,8 +106,8 @@ async def test_before_model_callback():
), ),
) )
runner = utils.TestInMemoryRunner(agent) runner = testing_utils.TestInMemoryRunner(agent)
assert utils.simplify_events( assert testing_utils.simplify_events(
await runner.run_async_with_new_session('test') await runner.run_async_with_new_session('test')
) == [ ) == [
('root_agent', 'before_model_callback'), ('root_agent', 'before_model_callback'),
@ -117,15 +117,15 @@ async def test_before_model_callback():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_before_model_callback_noop(): async def test_before_model_callback_noop():
responses = ['model_response'] responses = ['model_response']
mock_model = utils.MockModel.create(responses=responses) mock_model = testing_utils.MockModel.create(responses=responses)
agent = Agent( agent = Agent(
name='root_agent', name='root_agent',
model=mock_model, model=mock_model,
before_model_callback=noop_callback, before_model_callback=noop_callback,
) )
runner = utils.TestInMemoryRunner(agent) runner = testing_utils.TestInMemoryRunner(agent)
assert utils.simplify_events( assert testing_utils.simplify_events(
await runner.run_async_with_new_session('test') await runner.run_async_with_new_session('test')
) == [ ) == [
('root_agent', 'model_response'), ('root_agent', 'model_response'),
@ -135,7 +135,7 @@ async def test_before_model_callback_noop():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_after_model_callback(): async def test_after_model_callback():
responses = ['model_response'] responses = ['model_response']
mock_model = utils.MockModel.create(responses=responses) mock_model = testing_utils.MockModel.create(responses=responses)
agent = Agent( agent = Agent(
name='root_agent', name='root_agent',
model=mock_model, model=mock_model,
@ -144,8 +144,8 @@ async def test_after_model_callback():
), ),
) )
runner = utils.TestInMemoryRunner(agent) runner = testing_utils.TestInMemoryRunner(agent)
assert utils.simplify_events( assert testing_utils.simplify_events(
await runner.run_async_with_new_session('test') await runner.run_async_with_new_session('test')
) == [ ) == [
('root_agent', 'after_model_callback'), ('root_agent', 'after_model_callback'),
@ -155,7 +155,7 @@ async def test_after_model_callback():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_async_before_model_callback(): async def test_async_before_model_callback():
responses = ['model_response'] responses = ['model_response']
mock_model = utils.MockModel.create(responses=responses) mock_model = testing_utils.MockModel.create(responses=responses)
agent = Agent( agent = Agent(
name='root_agent', name='root_agent',
model=mock_model, model=mock_model,
@ -164,8 +164,8 @@ async def test_async_before_model_callback():
), ),
) )
runner = utils.TestInMemoryRunner(agent) runner = testing_utils.TestInMemoryRunner(agent)
assert utils.simplify_events( assert testing_utils.simplify_events(
await runner.run_async_with_new_session('test') await runner.run_async_with_new_session('test')
) == [ ) == [
('root_agent', 'async_before_model_callback'), ('root_agent', 'async_before_model_callback'),
@ -175,15 +175,15 @@ async def test_async_before_model_callback():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_async_before_model_callback_noop(): async def test_async_before_model_callback_noop():
responses = ['model_response'] responses = ['model_response']
mock_model = utils.MockModel.create(responses=responses) mock_model = testing_utils.MockModel.create(responses=responses)
agent = Agent( agent = Agent(
name='root_agent', name='root_agent',
model=mock_model, model=mock_model,
before_model_callback=async_noop_callback, before_model_callback=async_noop_callback,
) )
runner = utils.TestInMemoryRunner(agent) runner = testing_utils.TestInMemoryRunner(agent)
assert utils.simplify_events( assert testing_utils.simplify_events(
await runner.run_async_with_new_session('test') await runner.run_async_with_new_session('test')
) == [ ) == [
('root_agent', 'model_response'), ('root_agent', 'model_response'),
@ -193,7 +193,7 @@ async def test_async_before_model_callback_noop():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_async_after_model_callback(): async def test_async_after_model_callback():
responses = ['model_response'] responses = ['model_response']
mock_model = utils.MockModel.create(responses=responses) mock_model = testing_utils.MockModel.create(responses=responses)
agent = Agent( agent = Agent(
name='root_agent', name='root_agent',
model=mock_model, model=mock_model,
@ -202,8 +202,8 @@ async def test_async_after_model_callback():
), ),
) )
runner = utils.TestInMemoryRunner(agent) runner = testing_utils.TestInMemoryRunner(agent)
assert utils.simplify_events( assert testing_utils.simplify_events(
await runner.run_async_with_new_session('test') await runner.run_async_with_new_session('test')
) == [ ) == [
('root_agent', 'async_after_model_callback'), ('root_agent', 'async_after_model_callback'),

View File

@ -27,7 +27,7 @@ from google.genai import types
from pydantic import BaseModel from pydantic import BaseModel
import pytest import pytest
from .. import utils from .. import testing_utils
class CallbackType(Enum): class CallbackType(Enum):
@ -42,7 +42,9 @@ async def mock_async_before_cb_side_effect(
): ):
if ret_value: if ret_value:
return LlmResponse( return LlmResponse(
content=utils.ModelContent([types.Part.from_text(text=ret_value)]) content=testing_utils.ModelContent(
[types.Part.from_text(text=ret_value)]
)
) )
return None return None
@ -54,7 +56,9 @@ def mock_sync_before_cb_side_effect(
): ):
if ret_value: if ret_value:
return LlmResponse( return LlmResponse(
content=utils.ModelContent([types.Part.from_text(text=ret_value)]) content=testing_utils.ModelContent(
[types.Part.from_text(text=ret_value)]
)
) )
return None return None
@ -66,7 +70,9 @@ async def mock_async_after_cb_side_effect(
): ):
if ret_value: if ret_value:
return LlmResponse( return LlmResponse(
content=utils.ModelContent([types.Part.from_text(text=ret_value)]) content=testing_utils.ModelContent(
[types.Part.from_text(text=ret_value)]
)
) )
return None return None
@ -78,7 +84,9 @@ def mock_sync_after_cb_side_effect(
): ):
if ret_value: if ret_value:
return LlmResponse( return LlmResponse(
content=utils.ModelContent([types.Part.from_text(text=ret_value)]) content=testing_utils.ModelContent(
[types.Part.from_text(text=ret_value)]
)
) )
return None return None
@ -129,7 +137,7 @@ async def test_before_model_callbacks_chain(
expected_calls: List[int], expected_calls: List[int],
): ):
responses = ["model_response"] responses = ["model_response"]
mock_model = utils.MockModel.create(responses=responses) mock_model = testing_utils.MockModel.create(responses=responses)
mock_cbs = [] mock_cbs = []
for response, callback_type in callbacks: for response, callback_type in callbacks:
@ -154,9 +162,9 @@ async def test_before_model_callbacks_chain(
before_model_callback=[mock_cb for mock_cb in mock_cbs], before_model_callback=[mock_cb for mock_cb in mock_cbs],
) )
runner = utils.TestInMemoryRunner(agent) runner = testing_utils.TestInMemoryRunner(agent)
result = await runner.run_async_with_new_session("test") result = await runner.run_async_with_new_session("test")
assert utils.simplify_events(result) == [ assert testing_utils.simplify_events(result) == [
("root_agent", expected_response), ("root_agent", expected_response),
] ]
@ -191,7 +199,7 @@ async def test_after_model_callbacks_chain(
expected_calls: List[int], expected_calls: List[int],
): ):
responses = ["model_response"] responses = ["model_response"]
mock_model = utils.MockModel.create(responses=responses) mock_model = testing_utils.MockModel.create(responses=responses)
mock_cbs = [] mock_cbs = []
for response, callback_type in callbacks: for response, callback_type in callbacks:
@ -216,9 +224,9 @@ async def test_after_model_callbacks_chain(
after_model_callback=[mock_cb for mock_cb in mock_cbs], after_model_callback=[mock_cb for mock_cb in mock_cbs],
) )
runner = utils.TestInMemoryRunner(agent) runner = testing_utils.TestInMemoryRunner(agent)
result = await runner.run_async_with_new_session("test") result = await runner.run_async_with_new_session("test")
assert utils.simplify_events(result) == [ assert testing_utils.simplify_events(result) == [
("root_agent", expected_response), ("root_agent", expected_response),
] ]

View File

@ -21,7 +21,7 @@ from google.adk.models.base_llm import LlmRequest
from google.genai import types from google.genai import types
import pytest import pytest
from ... import utils from ... import testing_utils
@pytest.mark.asyncio @pytest.mark.asyncio
@ -31,7 +31,7 @@ async def test_no_examples():
config=types.GenerateContentConfig(system_instruction=""), config=types.GenerateContentConfig(system_instruction=""),
) )
agent = Agent(model="gemini-1.5-flash", name="agent", examples=[]) agent = Agent(model="gemini-1.5-flash", name="agent", examples=[])
invocation_context = await utils.create_invocation_context( invocation_context = await testing_utils.create_invocation_context(
agent=agent, user_content="" agent=agent, user_content=""
) )
@ -69,7 +69,7 @@ async def test_agent_examples():
name="agent", name="agent",
examples=example_list, examples=example_list,
) )
invocation_context = await utils.create_invocation_context( invocation_context = await testing_utils.create_invocation_context(
agent=agent, user_content="test" agent=agent, user_content="test"
) )
@ -122,7 +122,7 @@ async def test_agent_base_example_provider():
name="agent", name="agent",
examples=provider, examples=provider,
) )
invocation_context = await utils.create_invocation_context( invocation_context = await testing_utils.create_invocation_context(
agent=agent, user_content="test" agent=agent, user_content="test"
) )

View File

@ -18,7 +18,7 @@ from google.adk.agents.sequential_agent import SequentialAgent
from google.adk.tools import exit_loop from google.adk.tools import exit_loop
from google.genai.types import Part from google.genai.types import Part
from ... import utils from ... import testing_utils
def transfer_call_part(agent_name: str) -> Part: def transfer_call_part(agent_name: str) -> Part:
@ -38,7 +38,7 @@ def test_auto_to_auto():
'response1', 'response1',
'response2', 'response2',
] ]
mockModel = utils.MockModel.create(responses=response) mockModel = testing_utils.MockModel.create(responses=response)
# root (auto) - sub_agent_1 (auto) # root (auto) - sub_agent_1 (auto)
sub_agent_1 = Agent(name='sub_agent_1', model=mockModel) sub_agent_1 = Agent(name='sub_agent_1', model=mockModel)
root_agent = Agent( root_agent = Agent(
@ -47,17 +47,17 @@ def test_auto_to_auto():
sub_agents=[sub_agent_1], sub_agents=[sub_agent_1],
) )
runner = utils.InMemoryRunner(root_agent) runner = testing_utils.InMemoryRunner(root_agent)
# Asserts the transfer. # Asserts the transfer.
assert utils.simplify_events(runner.run('test1')) == [ assert testing_utils.simplify_events(runner.run('test1')) == [
('root_agent', transfer_call_part('sub_agent_1')), ('root_agent', transfer_call_part('sub_agent_1')),
('root_agent', TRANSFER_RESPONSE_PART), ('root_agent', TRANSFER_RESPONSE_PART),
('sub_agent_1', 'response1'), ('sub_agent_1', 'response1'),
] ]
# sub_agent_1 should still be the current agent. # sub_agent_1 should still be the current agent.
assert utils.simplify_events(runner.run('test2')) == [ assert testing_utils.simplify_events(runner.run('test2')) == [
('sub_agent_1', 'response2'), ('sub_agent_1', 'response2'),
] ]
@ -68,7 +68,7 @@ def test_auto_to_single():
'response1', 'response1',
'response2', 'response2',
] ]
mockModel = utils.MockModel.create(responses=response) mockModel = testing_utils.MockModel.create(responses=response)
# root (auto) - sub_agent_1 (single) # root (auto) - sub_agent_1 (single)
sub_agent_1 = Agent( sub_agent_1 = Agent(
name='sub_agent_1', name='sub_agent_1',
@ -80,17 +80,17 @@ def test_auto_to_single():
name='root_agent', model=mockModel, sub_agents=[sub_agent_1] name='root_agent', model=mockModel, sub_agents=[sub_agent_1]
) )
runner = utils.InMemoryRunner(root_agent) runner = testing_utils.InMemoryRunner(root_agent)
# Asserts the responses. # Asserts the responses.
assert utils.simplify_events(runner.run('test1')) == [ assert testing_utils.simplify_events(runner.run('test1')) == [
('root_agent', transfer_call_part('sub_agent_1')), ('root_agent', transfer_call_part('sub_agent_1')),
('root_agent', TRANSFER_RESPONSE_PART), ('root_agent', TRANSFER_RESPONSE_PART),
('sub_agent_1', 'response1'), ('sub_agent_1', 'response1'),
] ]
# root_agent should still be the current agent, becaues sub_agent_1 is single. # root_agent should still be the current agent, becaues sub_agent_1 is single.
assert utils.simplify_events(runner.run('test2')) == [ assert testing_utils.simplify_events(runner.run('test2')) == [
('root_agent', 'response2'), ('root_agent', 'response2'),
] ]
@ -103,7 +103,7 @@ def test_auto_to_auto_to_single():
'response1', 'response1',
'response2', 'response2',
] ]
mockModel = utils.MockModel.create(responses=response) mockModel = testing_utils.MockModel.create(responses=response)
# root (auto) - sub_agent_1 (auto) - sub_agent_1_1 (single) # root (auto) - sub_agent_1 (auto) - sub_agent_1_1 (single)
sub_agent_1_1 = Agent( sub_agent_1_1 = Agent(
name='sub_agent_1_1', name='sub_agent_1_1',
@ -118,10 +118,10 @@ def test_auto_to_auto_to_single():
name='root_agent', model=mockModel, sub_agents=[sub_agent_1] name='root_agent', model=mockModel, sub_agents=[sub_agent_1]
) )
runner = utils.InMemoryRunner(root_agent) runner = testing_utils.InMemoryRunner(root_agent)
# Asserts the responses. # Asserts the responses.
assert utils.simplify_events(runner.run('test1')) == [ assert testing_utils.simplify_events(runner.run('test1')) == [
('root_agent', transfer_call_part('sub_agent_1')), ('root_agent', transfer_call_part('sub_agent_1')),
('root_agent', TRANSFER_RESPONSE_PART), ('root_agent', TRANSFER_RESPONSE_PART),
('sub_agent_1', transfer_call_part('sub_agent_1_1')), ('sub_agent_1', transfer_call_part('sub_agent_1_1')),
@ -132,7 +132,7 @@ def test_auto_to_auto_to_single():
# sub_agent_1 should still be the current agent. sub_agent_1_1 is single so it should # sub_agent_1 should still be the current agent. sub_agent_1_1 is single so it should
# not be the current agent, otherwise the conversation will be tied to # not be the current agent, otherwise the conversation will be tied to
# sub_agent_1_1 forever. # sub_agent_1_1 forever.
assert utils.simplify_events(runner.run('test2')) == [ assert testing_utils.simplify_events(runner.run('test2')) == [
('sub_agent_1', 'response2'), ('sub_agent_1', 'response2'),
] ]
@ -145,7 +145,7 @@ def test_auto_to_sequential():
'response2', 'response2',
'response3', 'response3',
] ]
mockModel = utils.MockModel.create(responses=response) mockModel = testing_utils.MockModel.create(responses=response)
# root (auto) - sub_agent_1 (sequential) - sub_agent_1_1 (single) # root (auto) - sub_agent_1 (sequential) - sub_agent_1_1 (single)
# \ sub_agent_1_2 (single) # \ sub_agent_1_2 (single)
sub_agent_1_1 = Agent( sub_agent_1_1 = Agent(
@ -170,10 +170,10 @@ def test_auto_to_sequential():
sub_agents=[sub_agent_1], sub_agents=[sub_agent_1],
) )
runner = utils.InMemoryRunner(root_agent) runner = testing_utils.InMemoryRunner(root_agent)
# Asserts the transfer. # Asserts the transfer.
assert utils.simplify_events(runner.run('test1')) == [ assert testing_utils.simplify_events(runner.run('test1')) == [
('root_agent', transfer_call_part('sub_agent_1')), ('root_agent', transfer_call_part('sub_agent_1')),
('root_agent', TRANSFER_RESPONSE_PART), ('root_agent', TRANSFER_RESPONSE_PART),
('sub_agent_1_1', 'response1'), ('sub_agent_1_1', 'response1'),
@ -181,7 +181,7 @@ def test_auto_to_sequential():
] ]
# root_agent should still be the current agent because sub_agent_1 is sequential. # root_agent should still be the current agent because sub_agent_1 is sequential.
assert utils.simplify_events(runner.run('test2')) == [ assert testing_utils.simplify_events(runner.run('test2')) == [
('root_agent', 'response3'), ('root_agent', 'response3'),
] ]
@ -196,7 +196,7 @@ def test_auto_to_sequential_to_auto():
'response3', 'response3',
'response4', 'response4',
] ]
mockModel = utils.MockModel.create(responses=response) mockModel = testing_utils.MockModel.create(responses=response)
# root (auto) - sub_agent_1 (seq) - sub_agent_1_1 (single) # root (auto) - sub_agent_1 (seq) - sub_agent_1_1 (single)
# \ sub_agent_1_2 (auto) - sub_agent_1_2_1 (auto) # \ sub_agent_1_2 (auto) - sub_agent_1_2_1 (auto)
# \ sub_agent_1_3 (single) # \ sub_agent_1_3 (single)
@ -228,10 +228,10 @@ def test_auto_to_sequential_to_auto():
sub_agents=[sub_agent_1], sub_agents=[sub_agent_1],
) )
runner = utils.InMemoryRunner(root_agent) runner = testing_utils.InMemoryRunner(root_agent)
# Asserts the transfer. # Asserts the transfer.
assert utils.simplify_events(runner.run('test1')) == [ assert testing_utils.simplify_events(runner.run('test1')) == [
('root_agent', transfer_call_part('sub_agent_1')), ('root_agent', transfer_call_part('sub_agent_1')),
('root_agent', TRANSFER_RESPONSE_PART), ('root_agent', TRANSFER_RESPONSE_PART),
('sub_agent_1_1', 'response1'), ('sub_agent_1_1', 'response1'),
@ -242,7 +242,7 @@ def test_auto_to_sequential_to_auto():
] ]
# root_agent should still be the current agent because sub_agent_1 is sequential. # root_agent should still be the current agent because sub_agent_1 is sequential.
assert utils.simplify_events(runner.run('test2')) == [ assert testing_utils.simplify_events(runner.run('test2')) == [
('root_agent', 'response4'), ('root_agent', 'response4'),
] ]
@ -258,7 +258,7 @@ def test_auto_to_loop():
'response4', 'response4',
'response5', 'response5',
] ]
mockModel = utils.MockModel.create(responses=response) mockModel = testing_utils.MockModel.create(responses=response)
# root (auto) - sub_agent_1 (loop) - sub_agent_1_1 (single) # root (auto) - sub_agent_1 (loop) - sub_agent_1_1 (single)
# \ sub_agent_1_2 (single) # \ sub_agent_1_2 (single)
sub_agent_1_1 = Agent( sub_agent_1_1 = Agent(
@ -284,10 +284,10 @@ def test_auto_to_loop():
sub_agents=[sub_agent_1], sub_agents=[sub_agent_1],
) )
runner = utils.InMemoryRunner(root_agent) runner = testing_utils.InMemoryRunner(root_agent)
# Asserts the transfer. # Asserts the transfer.
assert utils.simplify_events(runner.run('test1')) == [ assert testing_utils.simplify_events(runner.run('test1')) == [
# Transfers to sub_agent_1. # Transfers to sub_agent_1.
('root_agent', transfer_call_part('sub_agent_1')), ('root_agent', transfer_call_part('sub_agent_1')),
('root_agent', TRANSFER_RESPONSE_PART), ('root_agent', TRANSFER_RESPONSE_PART),
@ -306,6 +306,6 @@ def test_auto_to_loop():
] ]
# root_agent should still be the current agent because sub_agent_1 is loop. # root_agent should still be the current agent because sub_agent_1 is loop.
assert utils.simplify_events(runner.run('test2')) == [ assert testing_utils.simplify_events(runner.run('test2')) == [
('root_agent', 'response5'), ('root_agent', 'response5'),
] ]

View File

@ -29,7 +29,7 @@ from google.adk.tools.tool_context import ToolContext
from google.genai import types from google.genai import types
import pytest import pytest
from ... import utils from ... import testing_utils
class CallbackType(Enum): class CallbackType(Enum):
@ -73,7 +73,7 @@ async def invoke_tool_with_callbacks(
return {"initial": "response"} return {"initial": "response"}
tool = FunctionTool(simple_fn) tool = FunctionTool(simple_fn)
model = utils.MockModel.create(responses=[]) model = testing_utils.MockModel.create(responses=[])
agent = Agent( agent = Agent(
name="agent", name="agent",
model=model, model=model,
@ -81,7 +81,7 @@ async def invoke_tool_with_callbacks(
before_tool_callback=before_cb, before_tool_callback=before_cb,
after_tool_callback=after_cb, after_tool_callback=after_cb,
) )
invocation_context = await utils.create_invocation_context( invocation_context = await testing_utils.create_invocation_context(
agent=agent, user_content="" agent=agent, user_content=""
) )
# Build function call event # Build function call event

View File

@ -17,7 +17,7 @@ from google.adk.tools import ToolContext
from google.adk.tools.long_running_tool import LongRunningFunctionTool from google.adk.tools.long_running_tool import LongRunningFunctionTool
from google.genai.types import Part from google.genai.types import Part
from ... import utils from ... import testing_utils
def test_async_function(): def test_async_function():
@ -28,7 +28,7 @@ def test_async_function():
'response3', 'response3',
'response4', 'response4',
] ]
mockModel = utils.MockModel.create(responses=responses) mockModel = testing_utils.MockModel.create(responses=responses)
function_called = 0 function_called = 0
def increase_by_one(x: int, tool_context: ToolContext) -> int: def increase_by_one(x: int, tool_context: ToolContext) -> int:
@ -43,14 +43,14 @@ def test_async_function():
model=mockModel, model=mockModel,
tools=[LongRunningFunctionTool(func=increase_by_one)], tools=[LongRunningFunctionTool(func=increase_by_one)],
) )
runner = utils.InMemoryRunner(agent) runner = testing_utils.InMemoryRunner(agent)
events = runner.run('test1') events = runner.run('test1')
# Asserts the requests. # Asserts the requests.
assert len(mockModel.requests) == 2 assert len(mockModel.requests) == 2
# 1 item: user content # 1 item: user content
assert mockModel.requests[0].contents == [ assert mockModel.requests[0].contents == [
utils.UserContent('test1'), testing_utils.UserContent('test1'),
] ]
increase_by_one_call = Part.from_function_call( increase_by_one_call = Part.from_function_call(
name='increase_by_one', args={'x': 1} name='increase_by_one', args={'x': 1}
@ -59,7 +59,7 @@ def test_async_function():
name='increase_by_one', response={'status': 'pending'} name='increase_by_one', response={'status': 'pending'}
) )
assert utils.simplify_contents(mockModel.requests[1].contents) == [ assert testing_utils.simplify_contents(mockModel.requests[1].contents) == [
('user', 'test1'), ('user', 'test1'),
('model', increase_by_one_call), ('model', increase_by_one_call),
('user', pending_response), ('user', pending_response),
@ -69,7 +69,7 @@ def test_async_function():
assert function_called == 1 assert function_called == 1
# Asserts the responses. # Asserts the responses.
assert utils.simplify_events(events) == [ assert testing_utils.simplify_events(events) == [
( (
'root_agent', 'root_agent',
Part.from_function_call(name='increase_by_one', args={'x': 1}), Part.from_function_call(name='increase_by_one', args={'x': 1}),
@ -88,45 +88,45 @@ def test_async_function():
still_waiting_response = Part.from_function_response( still_waiting_response = Part.from_function_response(
name='increase_by_one', response={'status': 'still waiting'} name='increase_by_one', response={'status': 'still waiting'}
) )
events = runner.run(utils.UserContent(still_waiting_response)) events = runner.run(testing_utils.UserContent(still_waiting_response))
# We have one new request. # We have one new request.
assert len(mockModel.requests) == 3 assert len(mockModel.requests) == 3
assert utils.simplify_contents(mockModel.requests[2].contents) == [ assert testing_utils.simplify_contents(mockModel.requests[2].contents) == [
('user', 'test1'), ('user', 'test1'),
('model', increase_by_one_call), ('model', increase_by_one_call),
('user', still_waiting_response), ('user', still_waiting_response),
] ]
assert utils.simplify_events(events) == [('root_agent', 'response2')] assert testing_utils.simplify_events(events) == [('root_agent', 'response2')]
# Calls when the result is ready. # Calls when the result is ready.
result_response = Part.from_function_response( result_response = Part.from_function_response(
name='increase_by_one', response={'result': 2} name='increase_by_one', response={'result': 2}
) )
events = runner.run(utils.UserContent(result_response)) events = runner.run(testing_utils.UserContent(result_response))
# We have one new request. # We have one new request.
assert len(mockModel.requests) == 4 assert len(mockModel.requests) == 4
assert utils.simplify_contents(mockModel.requests[3].contents) == [ assert testing_utils.simplify_contents(mockModel.requests[3].contents) == [
('user', 'test1'), ('user', 'test1'),
('model', increase_by_one_call), ('model', increase_by_one_call),
('user', result_response), ('user', result_response),
] ]
assert utils.simplify_events(events) == [('root_agent', 'response3')] assert testing_utils.simplify_events(events) == [('root_agent', 'response3')]
# Calls when the result is ready. Here we still accept the result and do # Calls when the result is ready. Here we still accept the result and do
# another summarization. Whether this is the right behavior is TBD. # another summarization. Whether this is the right behavior is TBD.
another_result_response = Part.from_function_response( another_result_response = Part.from_function_response(
name='increase_by_one', response={'result': 3} name='increase_by_one', response={'result': 3}
) )
events = runner.run(utils.UserContent(another_result_response)) events = runner.run(testing_utils.UserContent(another_result_response))
# We have one new request. # We have one new request.
assert len(mockModel.requests) == 5 assert len(mockModel.requests) == 5
assert utils.simplify_contents(mockModel.requests[4].contents) == [ assert testing_utils.simplify_contents(mockModel.requests[4].contents) == [
('user', 'test1'), ('user', 'test1'),
('model', increase_by_one_call), ('model', increase_by_one_call),
('user', another_result_response), ('user', another_result_response),
] ]
assert utils.simplify_events(events) == [('root_agent', 'response4')] assert testing_utils.simplify_events(events) == [('root_agent', 'response4')]
# At the end, function_called should still be 1. # At the end, function_called should still be 1.
assert function_called == 1 assert function_called == 1
@ -140,7 +140,7 @@ def test_async_function_with_none_response():
'response3', 'response3',
'response4', 'response4',
] ]
mockModel = utils.MockModel.create(responses=responses) mockModel = testing_utils.MockModel.create(responses=responses)
function_called = 0 function_called = 0
def increase_by_one(x: int, tool_context: ToolContext) -> int: def increase_by_one(x: int, tool_context: ToolContext) -> int:
@ -154,20 +154,20 @@ def test_async_function_with_none_response():
model=mockModel, model=mockModel,
tools=[LongRunningFunctionTool(func=increase_by_one)], tools=[LongRunningFunctionTool(func=increase_by_one)],
) )
runner = utils.InMemoryRunner(agent) runner = testing_utils.InMemoryRunner(agent)
events = runner.run('test1') events = runner.run('test1')
# Asserts the requests. # Asserts the requests.
assert len(mockModel.requests) == 2 assert len(mockModel.requests) == 2
# 1 item: user content # 1 item: user content
assert mockModel.requests[0].contents == [ assert mockModel.requests[0].contents == [
utils.UserContent('test1'), testing_utils.UserContent('test1'),
] ]
increase_by_one_call = Part.from_function_call( increase_by_one_call = Part.from_function_call(
name='increase_by_one', args={'x': 1} name='increase_by_one', args={'x': 1}
) )
assert utils.simplify_contents(mockModel.requests[1].contents) == [ assert testing_utils.simplify_contents(mockModel.requests[1].contents) == [
('user', 'test1'), ('user', 'test1'),
('model', increase_by_one_call), ('model', increase_by_one_call),
( (
@ -182,7 +182,7 @@ def test_async_function_with_none_response():
assert function_called == 1 assert function_called == 1
# Asserts the responses. # Asserts the responses.
assert utils.simplify_events(events) == [ assert testing_utils.simplify_events(events) == [
( (
'root_agent', 'root_agent',
Part.from_function_call(name='increase_by_one', args={'x': 1}), Part.from_function_call(name='increase_by_one', args={'x': 1}),
@ -200,45 +200,45 @@ def test_async_function_with_none_response():
still_waiting_response = Part.from_function_response( still_waiting_response = Part.from_function_response(
name='increase_by_one', response={'status': 'still waiting'} name='increase_by_one', response={'status': 'still waiting'}
) )
events = runner.run(utils.UserContent(still_waiting_response)) events = runner.run(testing_utils.UserContent(still_waiting_response))
# We have one new request. # We have one new request.
assert len(mockModel.requests) == 3 assert len(mockModel.requests) == 3
assert utils.simplify_contents(mockModel.requests[2].contents) == [ assert testing_utils.simplify_contents(mockModel.requests[2].contents) == [
('user', 'test1'), ('user', 'test1'),
('model', increase_by_one_call), ('model', increase_by_one_call),
('user', still_waiting_response), ('user', still_waiting_response),
] ]
assert utils.simplify_events(events) == [('root_agent', 'response2')] assert testing_utils.simplify_events(events) == [('root_agent', 'response2')]
# Calls when the result is ready. # Calls when the result is ready.
result_response = Part.from_function_response( result_response = Part.from_function_response(
name='increase_by_one', response={'result': 2} name='increase_by_one', response={'result': 2}
) )
events = runner.run(utils.UserContent(result_response)) events = runner.run(testing_utils.UserContent(result_response))
# We have one new request. # We have one new request.
assert len(mockModel.requests) == 4 assert len(mockModel.requests) == 4
assert utils.simplify_contents(mockModel.requests[3].contents) == [ assert testing_utils.simplify_contents(mockModel.requests[3].contents) == [
('user', 'test1'), ('user', 'test1'),
('model', increase_by_one_call), ('model', increase_by_one_call),
('user', result_response), ('user', result_response),
] ]
assert utils.simplify_events(events) == [('root_agent', 'response3')] assert testing_utils.simplify_events(events) == [('root_agent', 'response3')]
# Calls when the result is ready. Here we still accept the result and do # Calls when the result is ready. Here we still accept the result and do
# another summarization. Whether this is the right behavior is TBD. # another summarization. Whether this is the right behavior is TBD.
another_result_response = Part.from_function_response( another_result_response = Part.from_function_response(
name='increase_by_one', response={'result': 3} name='increase_by_one', response={'result': 3}
) )
events = runner.run(utils.UserContent(another_result_response)) events = runner.run(testing_utils.UserContent(another_result_response))
# We have one new request. # We have one new request.
assert len(mockModel.requests) == 5 assert len(mockModel.requests) == 5
assert utils.simplify_contents(mockModel.requests[4].contents) == [ assert testing_utils.simplify_contents(mockModel.requests[4].contents) == [
('user', 'test1'), ('user', 'test1'),
('model', increase_by_one_call), ('model', increase_by_one_call),
('user', another_result_response), ('user', another_result_response),
] ]
assert utils.simplify_events(events) == [('root_agent', 'response4')] assert testing_utils.simplify_events(events) == [('root_agent', 'response4')]
# At the end, function_called should still be 1. # At the end, function_called should still be 1.
assert function_called == 1 assert function_called == 1

View File

@ -28,7 +28,7 @@ from google.adk.tools import AuthToolArguments
from google.adk.tools import ToolContext from google.adk.tools import ToolContext
from google.genai import types from google.genai import types
from ... import utils from ... import testing_utils
def function_call(function_call_id, name, args: dict[str, Any]) -> types.Part: def function_call(function_call_id, name, args: dict[str, Any]) -> types.Part:
@ -95,7 +95,7 @@ def test_function_request_euc():
), ),
) )
mock_model = utils.MockModel.create(responses=responses) mock_model = testing_utils.MockModel.create(responses=responses)
def call_external_api1(tool_context: ToolContext) -> Optional[int]: def call_external_api1(tool_context: ToolContext) -> Optional[int]:
tool_context.request_credential(auth_config1) tool_context.request_credential(auth_config1)
@ -108,7 +108,7 @@ def test_function_request_euc():
model=mock_model, model=mock_model,
tools=[call_external_api1, call_external_api2], tools=[call_external_api1, call_external_api2],
) )
runner = utils.InMemoryRunner(agent) runner = testing_utils.InMemoryRunner(agent)
events = runner.run('test') events = runner.run('test')
assert events[0].content.parts[0].function_call is not None assert events[0].content.parts[0].function_call is not None
assert events[0].content.parts[1].function_call is not None assert events[0].content.parts[1].function_call is not None
@ -169,7 +169,7 @@ def test_function_get_auth_response():
], ],
] ]
mock_model = utils.MockModel.create(responses=responses) mock_model = testing_utils.MockModel.create(responses=responses)
function_invoked = 0 function_invoked = 0
auth_config1 = AuthConfig( auth_config1 = AuthConfig(
@ -307,7 +307,7 @@ def test_function_get_auth_response():
model=mock_model, model=mock_model,
tools=[call_external_api1, call_external_api2], tools=[call_external_api1, call_external_api2],
) )
runner = utils.InMemoryRunner(agent) runner = testing_utils.InMemoryRunner(agent)
runner.run('test') runner.run('test')
request_euc_function_call_event = runner.session.events[-3] request_euc_function_call_event = runner.session.events[-3]
function_response1 = types.FunctionResponse( function_response1 = types.FunctionResponse(

View File

@ -17,7 +17,7 @@ from typing import Any
from google.adk.agents import Agent from google.adk.agents import Agent
from google.genai import types from google.genai import types
from ... import utils from ... import testing_utils
def function_call(args: dict[str, Any]) -> types.Part: def function_call(args: dict[str, Any]) -> types.Part:
@ -37,7 +37,7 @@ def test_sequential_calls():
function_call({'x': 3}), function_call({'x': 3}),
'response1', 'response1',
] ]
mockModel = utils.MockModel.create(responses=responses) mockModel = testing_utils.MockModel.create(responses=responses)
function_called = 0 function_called = 0
def increase_by_one(x: int) -> int: def increase_by_one(x: int) -> int:
@ -46,8 +46,8 @@ def test_sequential_calls():
return x + 1 return x + 1
agent = Agent(name='root_agent', model=mockModel, tools=[increase_by_one]) agent = Agent(name='root_agent', model=mockModel, tools=[increase_by_one])
runner = utils.InMemoryRunner(agent) runner = testing_utils.InMemoryRunner(agent)
result = utils.simplify_events(runner.run('test')) result = testing_utils.simplify_events(runner.run('test'))
assert result == [ assert result == [
('root_agent', function_call({'x': 1})), ('root_agent', function_call({'x': 1})),
('root_agent', function_response({'result': 2})), ('root_agent', function_response({'result': 2})),
@ -61,17 +61,17 @@ def test_sequential_calls():
# Asserts the requests. # Asserts the requests.
assert len(mockModel.requests) == 4 assert len(mockModel.requests) == 4
# 1 item: user content # 1 item: user content
assert utils.simplify_contents(mockModel.requests[0].contents) == [ assert testing_utils.simplify_contents(mockModel.requests[0].contents) == [
('user', 'test') ('user', 'test')
] ]
# 3 items: user content, functaion call / response for the 1st call # 3 items: user content, functaion call / response for the 1st call
assert utils.simplify_contents(mockModel.requests[1].contents) == [ assert testing_utils.simplify_contents(mockModel.requests[1].contents) == [
('user', 'test'), ('user', 'test'),
('model', function_call({'x': 1})), ('model', function_call({'x': 1})),
('user', function_response({'result': 2})), ('user', function_response({'result': 2})),
] ]
# 5 items: user content, functaion call / response for two calls # 5 items: user content, functaion call / response for two calls
assert utils.simplify_contents(mockModel.requests[2].contents) == [ assert testing_utils.simplify_contents(mockModel.requests[2].contents) == [
('user', 'test'), ('user', 'test'),
('model', function_call({'x': 1})), ('model', function_call({'x': 1})),
('user', function_response({'result': 2})), ('user', function_response({'result': 2})),
@ -79,7 +79,7 @@ def test_sequential_calls():
('user', function_response({'result': 3})), ('user', function_response({'result': 3})),
] ]
# 7 items: user content, functaion call / response for three calls # 7 items: user content, functaion call / response for three calls
assert utils.simplify_contents(mockModel.requests[3].contents) == [ assert testing_utils.simplify_contents(mockModel.requests[3].contents) == [
('user', 'test'), ('user', 'test'),
('model', function_call({'x': 1})), ('model', function_call({'x': 1})),
('user', function_response({'result': 2})), ('user', function_response({'result': 2})),

View File

@ -22,7 +22,7 @@ from google.adk.tools.function_tool import FunctionTool
from google.genai import types from google.genai import types
import pytest import pytest
from ... import utils from ... import testing_utils
def test_simple_function(): def test_simple_function():
@ -40,7 +40,7 @@ def test_simple_function():
'response4', 'response4',
] ]
function_called = 0 function_called = 0
mock_model = utils.MockModel.create(responses=responses) mock_model = testing_utils.MockModel.create(responses=responses)
def increase_by_one(x: int) -> int: def increase_by_one(x: int) -> int:
nonlocal function_called nonlocal function_called
@ -48,18 +48,18 @@ def test_simple_function():
return x + 1 return x + 1
agent = Agent(name='root_agent', model=mock_model, tools=[increase_by_one]) agent = Agent(name='root_agent', model=mock_model, tools=[increase_by_one])
runner = utils.InMemoryRunner(agent) runner = testing_utils.InMemoryRunner(agent)
assert utils.simplify_events(runner.run('test')) == [ assert testing_utils.simplify_events(runner.run('test')) == [
('root_agent', function_call_1), ('root_agent', function_call_1),
('root_agent', function_respones_2), ('root_agent', function_respones_2),
('root_agent', 'response1'), ('root_agent', 'response1'),
] ]
# Asserts the requests. # Asserts the requests.
assert utils.simplify_contents(mock_model.requests[0].contents) == [ assert testing_utils.simplify_contents(mock_model.requests[0].contents) == [
('user', 'test') ('user', 'test')
] ]
assert utils.simplify_contents(mock_model.requests[1].contents) == [ assert testing_utils.simplify_contents(mock_model.requests[1].contents) == [
('user', 'test'), ('user', 'test'),
('model', function_call_1), ('model', function_call_1),
('user', function_respones_2), ('user', function_respones_2),
@ -96,7 +96,7 @@ async def test_async_function():
'response4', 'response4',
] ]
function_called = 0 function_called = 0
mock_model = utils.MockModel.create(responses=responses) mock_model = testing_utils.MockModel.create(responses=responses)
async def increase_by_one(x: int) -> int: async def increase_by_one(x: int) -> int:
nonlocal function_called nonlocal function_called
@ -118,19 +118,19 @@ async def test_async_function():
model=mock_model, model=mock_model,
tools=[increase_by_one, multiple_by_two, multiple_by_two_sync], tools=[increase_by_one, multiple_by_two, multiple_by_two_sync],
) )
runner = utils.TestInMemoryRunner(agent) runner = testing_utils.TestInMemoryRunner(agent)
events = await runner.run_async_with_new_session('test') events = await runner.run_async_with_new_session('test')
assert utils.simplify_events(events) == [ assert testing_utils.simplify_events(events) == [
('root_agent', function_calls), ('root_agent', function_calls),
('root_agent', function_responses), ('root_agent', function_responses),
('root_agent', 'response1'), ('root_agent', 'response1'),
] ]
# Asserts the requests. # Asserts the requests.
assert utils.simplify_contents(mock_model.requests[0].contents) == [ assert testing_utils.simplify_contents(mock_model.requests[0].contents) == [
('user', 'test') ('user', 'test')
] ]
assert utils.simplify_contents(mock_model.requests[1].contents) == [ assert testing_utils.simplify_contents(mock_model.requests[1].contents) == [
('user', 'test'), ('user', 'test'),
('model', function_calls), ('model', function_calls),
('user', function_responses), ('user', function_responses),
@ -167,7 +167,7 @@ async def test_function_tool():
'response4', 'response4',
] ]
function_called = 0 function_called = 0
mock_model = utils.MockModel.create(responses=responses) mock_model = testing_utils.MockModel.create(responses=responses)
async def increase_by_one(x: int) -> int: async def increase_by_one(x: int) -> int:
nonlocal function_called nonlocal function_called
@ -195,19 +195,19 @@ async def test_function_tool():
model=mock_model, model=mock_model,
tools=[wrapped_increase_by_one, multiple_by_two, multiple_by_two_sync], tools=[wrapped_increase_by_one, multiple_by_two, multiple_by_two_sync],
) )
runner = utils.TestInMemoryRunner(agent) runner = testing_utils.TestInMemoryRunner(agent)
events = await runner.run_async_with_new_session('test') events = await runner.run_async_with_new_session('test')
assert utils.simplify_events(events) == [ assert testing_utils.simplify_events(events) == [
('root_agent', function_calls), ('root_agent', function_calls),
('root_agent', function_responses), ('root_agent', function_responses),
('root_agent', 'response1'), ('root_agent', 'response1'),
] ]
# Asserts the requests. # Asserts the requests.
assert utils.simplify_contents(mock_model.requests[0].contents) == [ assert testing_utils.simplify_contents(mock_model.requests[0].contents) == [
('user', 'test') ('user', 'test')
] ]
assert utils.simplify_contents(mock_model.requests[1].contents) == [ assert testing_utils.simplify_contents(mock_model.requests[1].contents) == [
('user', 'test'), ('user', 'test'),
('model', function_calls), ('model', function_calls),
('user', function_responses), ('user', function_responses),
@ -218,7 +218,7 @@ async def test_function_tool():
def test_update_state(): def test_update_state():
mock_model = utils.MockModel.create( mock_model = testing_utils.MockModel.create(
responses=[ responses=[
types.Part.from_function_call(name='update_state', args={}), types.Part.from_function_call(name='update_state', args={}),
'response1', 'response1',
@ -229,7 +229,7 @@ def test_update_state():
tool_context.state['x'] = 1 tool_context.state['x'] = 1
agent = Agent(name='root_agent', model=mock_model, tools=[update_state]) agent = Agent(name='root_agent', model=mock_model, tools=[update_state])
runner = utils.InMemoryRunner(agent) runner = testing_utils.InMemoryRunner(agent)
runner.run('test') runner.run('test')
assert runner.session.state['x'] == 1 assert runner.session.state['x'] == 1
@ -239,13 +239,13 @@ def test_function_call_id():
types.Part.from_function_call(name='increase_by_one', args={'x': 1}), types.Part.from_function_call(name='increase_by_one', args={'x': 1}),
'response1', 'response1',
] ]
mock_model = utils.MockModel.create(responses=responses) mock_model = testing_utils.MockModel.create(responses=responses)
def increase_by_one(x: int) -> int: def increase_by_one(x: int) -> int:
return x + 1 return x + 1
agent = Agent(name='root_agent', model=mock_model, tools=[increase_by_one]) agent = Agent(name='root_agent', model=mock_model, tools=[increase_by_one])
runner = utils.InMemoryRunner(agent) runner = testing_utils.InMemoryRunner(agent)
events = runner.run('test') events = runner.run('test')
for reqeust in mock_model.requests: for reqeust in mock_model.requests:
for content in reqeust.contents: for content in reqeust.contents:

View File

@ -18,7 +18,7 @@ from google.adk.models import LlmRequest
from google.genai import types from google.genai import types
import pytest import pytest
from ... import utils from ... import testing_utils
@pytest.mark.asyncio @pytest.mark.asyncio
@ -28,7 +28,9 @@ async def test_no_description():
config=types.GenerateContentConfig(system_instruction=""), config=types.GenerateContentConfig(system_instruction=""),
) )
agent = Agent(model="gemini-1.5-flash", name="agent") agent = Agent(model="gemini-1.5-flash", name="agent")
invocation_context = await utils.create_invocation_context(agent=agent) invocation_context = await testing_utils.create_invocation_context(
agent=agent
)
async for _ in identity.request_processor.run_async( async for _ in identity.request_processor.run_async(
invocation_context, invocation_context,
@ -52,7 +54,9 @@ async def test_with_description():
name="agent", name="agent",
description="test description", description="test description",
) )
invocation_context = await utils.create_invocation_context(agent=agent) invocation_context = await testing_utils.create_invocation_context(
agent=agent
)
async for _ in identity.request_processor.run_async( async for _ in identity.request_processor.run_async(
invocation_context, invocation_context,

View File

@ -20,7 +20,7 @@ from google.adk.sessions import Session
from google.genai import types from google.genai import types
import pytest import pytest
from ... import utils from ... import testing_utils
@pytest.mark.asyncio @pytest.mark.asyncio
@ -36,7 +36,9 @@ async def test_build_system_instruction():
{{customer_int }, { non-identifier-float}}, \ {{customer_int }, { non-identifier-float}}, \
{'key1': 'value1'} and {{'key2': 'value2'}}."""), {'key1': 'value1'} and {{'key2': 'value2'}}."""),
) )
invocation_context = await utils.create_invocation_context(agent=agent) invocation_context = await testing_utils.create_invocation_context(
agent=agent
)
invocation_context.session = Session( invocation_context.session = Session(
app_name="test_app", app_name="test_app",
user_id="test_user", user_id="test_user",
@ -75,7 +77,9 @@ async def test_function_system_instruction():
name="agent", name="agent",
instruction=build_function_instruction, instruction=build_function_instruction,
) )
invocation_context = await utils.create_invocation_context(agent=agent) invocation_context = await testing_utils.create_invocation_context(
agent=agent
)
invocation_context.session = Session( invocation_context.session = Session(
app_name="test_app", app_name="test_app",
user_id="test_user", user_id="test_user",
@ -117,7 +121,9 @@ async def test_async_function_system_instruction():
name="agent", name="agent",
instruction=build_function_instruction, instruction=build_function_instruction,
) )
invocation_context = await utils.create_invocation_context(agent=agent) invocation_context = await testing_utils.create_invocation_context(
agent=agent
)
invocation_context.session = Session( invocation_context.session = Session(
app_name="test_app", app_name="test_app",
user_id="test_user", user_id="test_user",
@ -156,7 +162,9 @@ async def test_global_system_instruction():
model="gemini-1.5-flash", model="gemini-1.5-flash",
config=types.GenerateContentConfig(system_instruction=""), config=types.GenerateContentConfig(system_instruction=""),
) )
invocation_context = await utils.create_invocation_context(agent=sub_agent) invocation_context = await testing_utils.create_invocation_context(
agent=sub_agent
)
invocation_context.session = Session( invocation_context.session = Session(
app_name="test_app", app_name="test_app",
user_id="test_user", user_id="test_user",
@ -198,7 +206,9 @@ async def test_function_global_system_instruction():
model="gemini-1.5-flash", model="gemini-1.5-flash",
config=types.GenerateContentConfig(system_instruction=""), config=types.GenerateContentConfig(system_instruction=""),
) )
invocation_context = await utils.create_invocation_context(agent=sub_agent) invocation_context = await testing_utils.create_invocation_context(
agent=sub_agent
)
invocation_context.session = Session( invocation_context.session = Session(
app_name="test_app", app_name="test_app",
user_id="test_user", user_id="test_user",
@ -240,7 +250,9 @@ async def test_async_function_global_system_instruction():
model="gemini-1.5-flash", model="gemini-1.5-flash",
config=types.GenerateContentConfig(system_instruction=""), config=types.GenerateContentConfig(system_instruction=""),
) )
invocation_context = await utils.create_invocation_context(agent=sub_agent) invocation_context = await testing_utils.create_invocation_context(
agent=sub_agent
)
invocation_context.session = Session( invocation_context.session = Session(
app_name="test_app", app_name="test_app",
user_id="test_user", user_id="test_user",
@ -272,7 +284,9 @@ async def test_build_system_instruction_with_namespace():
"""Use the echo_info tool to echo { customerId }, {app:key}, {user:key}, {a:key}.""" """Use the echo_info tool to echo { customerId }, {app:key}, {user:key}, {a:key}."""
), ),
) )
invocation_context = await utils.create_invocation_context(agent=agent) invocation_context = await testing_utils.create_invocation_context(
agent=agent
)
invocation_context.session = Session( invocation_context.session = Session(
app_name="test_app", app_name="test_app",
user_id="test_user", user_id="test_user",

View File

@ -23,7 +23,7 @@ from google.genai import types
from pydantic import BaseModel from pydantic import BaseModel
import pytest import pytest
from ... import utils from ... import testing_utils
class MockBeforeModelCallback(BaseModel): class MockBeforeModelCallback(BaseModel):
@ -35,7 +35,7 @@ class MockBeforeModelCallback(BaseModel):
llm_request: LlmRequest, llm_request: LlmRequest,
) -> LlmResponse: ) -> LlmResponse:
return LlmResponse( return LlmResponse(
content=utils.ModelContent( content=testing_utils.ModelContent(
[types.Part.from_text(text=self.mock_response)] [types.Part.from_text(text=self.mock_response)]
) )
) )
@ -50,7 +50,7 @@ class MockAfterModelCallback(BaseModel):
llm_response: LlmResponse, llm_response: LlmResponse,
) -> LlmResponse: ) -> LlmResponse:
return LlmResponse( return LlmResponse(
content=utils.ModelContent( content=testing_utils.ModelContent(
[types.Part.from_text(text=self.mock_response)] [types.Part.from_text(text=self.mock_response)]
) )
) )
@ -62,7 +62,7 @@ def noop_callback(**kwargs) -> Optional[LlmResponse]:
def test_before_model_callback(): def test_before_model_callback():
responses = ['model_response'] responses = ['model_response']
mock_model = utils.MockModel.create(responses=responses) mock_model = testing_utils.MockModel.create(responses=responses)
agent = Agent( agent = Agent(
name='root_agent', name='root_agent',
model=mock_model, model=mock_model,
@ -71,30 +71,30 @@ def test_before_model_callback():
), ),
) )
runner = utils.InMemoryRunner(agent) runner = testing_utils.InMemoryRunner(agent)
assert utils.simplify_events(runner.run('test')) == [ assert testing_utils.simplify_events(runner.run('test')) == [
('root_agent', 'before_model_callback'), ('root_agent', 'before_model_callback'),
] ]
def test_before_model_callback_noop(): def test_before_model_callback_noop():
responses = ['model_response'] responses = ['model_response']
mock_model = utils.MockModel.create(responses=responses) mock_model = testing_utils.MockModel.create(responses=responses)
agent = Agent( agent = Agent(
name='root_agent', name='root_agent',
model=mock_model, model=mock_model,
before_model_callback=noop_callback, before_model_callback=noop_callback,
) )
runner = utils.InMemoryRunner(agent) runner = testing_utils.InMemoryRunner(agent)
assert utils.simplify_events(runner.run('test')) == [ assert testing_utils.simplify_events(runner.run('test')) == [
('root_agent', 'model_response'), ('root_agent', 'model_response'),
] ]
def test_before_model_callback_end(): def test_before_model_callback_end():
responses = ['model_response'] responses = ['model_response']
mock_model = utils.MockModel.create(responses=responses) mock_model = testing_utils.MockModel.create(responses=responses)
agent = Agent( agent = Agent(
name='root_agent', name='root_agent',
model=mock_model, model=mock_model,
@ -103,15 +103,15 @@ def test_before_model_callback_end():
), ),
) )
runner = utils.InMemoryRunner(agent) runner = testing_utils.InMemoryRunner(agent)
assert utils.simplify_events(runner.run('test')) == [ assert testing_utils.simplify_events(runner.run('test')) == [
('root_agent', 'before_model_callback'), ('root_agent', 'before_model_callback'),
] ]
def test_after_model_callback(): def test_after_model_callback():
responses = ['model_response'] responses = ['model_response']
mock_model = utils.MockModel.create(responses=responses) mock_model = testing_utils.MockModel.create(responses=responses)
agent = Agent( agent = Agent(
name='root_agent', name='root_agent',
model=mock_model, model=mock_model,
@ -120,8 +120,8 @@ def test_after_model_callback():
), ),
) )
runner = utils.InMemoryRunner(agent) runner = testing_utils.InMemoryRunner(agent)
assert utils.simplify_events(runner.run('test')) == [ assert testing_utils.simplify_events(runner.run('test')) == [
('root_agent', 'after_model_callback'), ('root_agent', 'after_model_callback'),
] ]
@ -129,14 +129,14 @@ def test_after_model_callback():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_after_model_callback_noop(): async def test_after_model_callback_noop():
responses = ['model_response'] responses = ['model_response']
mock_model = utils.MockModel.create(responses=responses) mock_model = testing_utils.MockModel.create(responses=responses)
agent = Agent( agent = Agent(
name='root_agent', name='root_agent',
model=mock_model, model=mock_model,
after_model_callback=noop_callback, after_model_callback=noop_callback,
) )
runner = utils.TestInMemoryRunner(agent) runner = testing_utils.TestInMemoryRunner(agent)
assert utils.simplify_events( assert testing_utils.simplify_events(
await runner.run_async_with_new_session('test') await runner.run_async_with_new_session('test')
) == [('root_agent', 'model_response')] ) == [('root_agent', 'model_response')]

View File

@ -17,7 +17,7 @@ from google.adk.tools import ToolContext
from google.genai.types import Part from google.genai.types import Part
from pydantic import BaseModel from pydantic import BaseModel
from ... import utils from ... import testing_utils
def test_output_schema(): def test_output_schema():
@ -27,7 +27,7 @@ def test_output_schema():
response = [ response = [
'response1', 'response1',
] ]
mockModel = utils.MockModel.create(responses=response) mockModel = testing_utils.MockModel.create(responses=response)
root_agent = Agent( root_agent = Agent(
name='root_agent', name='root_agent',
model=mockModel, model=mockModel,
@ -36,9 +36,9 @@ def test_output_schema():
disallow_transfer_to_peers=True, disallow_transfer_to_peers=True,
) )
runner = utils.InMemoryRunner(root_agent) runner = testing_utils.InMemoryRunner(root_agent)
assert utils.simplify_events(runner.run('test1')) == [ assert testing_utils.simplify_events(runner.run('test1')) == [
('root_agent', 'response1'), ('root_agent', 'response1'),
] ]
assert len(mockModel.requests) == 1 assert len(mockModel.requests) == 1

View File

@ -21,7 +21,7 @@ from google.genai import types
from google.genai.types import Part from google.genai.types import Part
from pydantic import BaseModel from pydantic import BaseModel
from ... import utils from ... import testing_utils
def simple_function(input_str: str) -> str: def simple_function(input_str: str) -> str:
@ -76,7 +76,7 @@ def test_before_tool_callback():
types.Part.from_function_call(name='simple_function', args={}), types.Part.from_function_call(name='simple_function', args={}),
'response1', 'response1',
] ]
mock_model = utils.MockModel.create(responses=responses) mock_model = testing_utils.MockModel.create(responses=responses)
agent = Agent( agent = Agent(
name='root_agent', name='root_agent',
model=mock_model, model=mock_model,
@ -86,8 +86,8 @@ def test_before_tool_callback():
tools=[simple_function], tools=[simple_function],
) )
runner = utils.InMemoryRunner(agent) runner = testing_utils.InMemoryRunner(agent)
assert utils.simplify_events(runner.run('test')) == [ assert testing_utils.simplify_events(runner.run('test')) == [
('root_agent', Part.from_function_call(name='simple_function', args={})), ('root_agent', Part.from_function_call(name='simple_function', args={})),
( (
'root_agent', 'root_agent',
@ -106,7 +106,7 @@ def test_before_tool_callback_noop():
), ),
'response1', 'response1',
] ]
mock_model = utils.MockModel.create(responses=responses) mock_model = testing_utils.MockModel.create(responses=responses)
agent = Agent( agent = Agent(
name='root_agent', name='root_agent',
model=mock_model, model=mock_model,
@ -114,8 +114,8 @@ def test_before_tool_callback_noop():
tools=[simple_function], tools=[simple_function],
) )
runner = utils.InMemoryRunner(agent) runner = testing_utils.InMemoryRunner(agent)
assert utils.simplify_events(runner.run('test')) == [ assert testing_utils.simplify_events(runner.run('test')) == [
( (
'root_agent', 'root_agent',
Part.from_function_call( Part.from_function_call(
@ -138,7 +138,7 @@ def test_before_tool_callback_modify_tool_request():
types.Part.from_function_call(name='simple_function', args={}), types.Part.from_function_call(name='simple_function', args={}),
'response1', 'response1',
] ]
mock_model = utils.MockModel.create(responses=responses) mock_model = testing_utils.MockModel.create(responses=responses)
agent = Agent( agent = Agent(
name='root_agent', name='root_agent',
model=mock_model, model=mock_model,
@ -149,8 +149,8 @@ def test_before_tool_callback_modify_tool_request():
tools=[simple_function], tools=[simple_function],
) )
runner = utils.InMemoryRunner(agent) runner = testing_utils.InMemoryRunner(agent)
assert utils.simplify_events(runner.run('test')) == [ assert testing_utils.simplify_events(runner.run('test')) == [
('root_agent', Part.from_function_call(name='simple_function', args={})), ('root_agent', Part.from_function_call(name='simple_function', args={})),
( (
'root_agent', 'root_agent',
@ -170,7 +170,7 @@ def test_after_tool_callback():
), ),
'response1', 'response1',
] ]
mock_model = utils.MockModel.create(responses=responses) mock_model = testing_utils.MockModel.create(responses=responses)
agent = Agent( agent = Agent(
name='root_agent', name='root_agent',
model=mock_model, model=mock_model,
@ -180,8 +180,8 @@ def test_after_tool_callback():
tools=[simple_function], tools=[simple_function],
) )
runner = utils.InMemoryRunner(agent) runner = testing_utils.InMemoryRunner(agent)
assert utils.simplify_events(runner.run('test')) == [ assert testing_utils.simplify_events(runner.run('test')) == [
( (
'root_agent', 'root_agent',
Part.from_function_call( Part.from_function_call(
@ -205,7 +205,7 @@ def test_after_tool_callback_noop():
), ),
'response1', 'response1',
] ]
mock_model = utils.MockModel.create(responses=responses) mock_model = testing_utils.MockModel.create(responses=responses)
agent = Agent( agent = Agent(
name='root_agent', name='root_agent',
model=mock_model, model=mock_model,
@ -213,8 +213,8 @@ def test_after_tool_callback_noop():
tools=[simple_function], tools=[simple_function],
) )
runner = utils.InMemoryRunner(agent) runner = testing_utils.InMemoryRunner(agent)
assert utils.simplify_events(runner.run('test')) == [ assert testing_utils.simplify_events(runner.run('test')) == [
( (
'root_agent', 'root_agent',
Part.from_function_call( Part.from_function_call(
@ -239,7 +239,7 @@ def test_after_tool_callback_modify_tool_response():
), ),
'response1', 'response1',
] ]
mock_model = utils.MockModel.create(responses=responses) mock_model = testing_utils.MockModel.create(responses=responses)
agent = Agent( agent = Agent(
name='root_agent', name='root_agent',
model=mock_model, model=mock_model,
@ -250,8 +250,8 @@ def test_after_tool_callback_modify_tool_response():
tools=[simple_function], tools=[simple_function],
) )
runner = utils.InMemoryRunner(agent) runner = testing_utils.InMemoryRunner(agent)
assert utils.simplify_events(runner.run('test')) == [ assert testing_utils.simplify_events(runner.run('test')) == [
( (
'root_agent', 'root_agent',
Part.from_function_call( Part.from_function_call(

View File

@ -18,7 +18,7 @@ from google.adk.models import LlmResponse
from google.genai import types from google.genai import types
import pytest import pytest
from .. import utils from .. import testing_utils
def test_streaming(): def test_streaming():
@ -26,7 +26,7 @@ def test_streaming():
turn_complete=True, turn_complete=True,
) )
mock_model = utils.MockModel.create([response1]) mock_model = testing_utils.MockModel.create([response1])
root_agent = Agent( root_agent = Agent(
name='root_agent', name='root_agent',
@ -34,7 +34,7 @@ def test_streaming():
tools=[], tools=[],
) )
runner = utils.InMemoryRunner( runner = testing_utils.InMemoryRunner(
root_agent=root_agent, response_modalities=['AUDIO'] root_agent=root_agent, response_modalities=['AUDIO']
) )
live_request_queue = LiveRequestQueue() live_request_queue = LiveRequestQueue()

View File

@ -17,7 +17,7 @@ from google.adk.tools.function_tool import FunctionTool
from google.adk.tools.retrieval.vertex_ai_rag_retrieval import VertexAiRagRetrieval from google.adk.tools.retrieval.vertex_ai_rag_retrieval import VertexAiRagRetrieval
from google.genai import types from google.genai import types
from ... import utils from ... import testing_utils
def noop_tool(x: str) -> str: def noop_tool(x: str) -> str:
@ -28,7 +28,7 @@ def test_vertex_rag_retrieval_for_gemini_1_x():
responses = [ responses = [
'response1', 'response1',
] ]
mockModel = utils.MockModel.create(responses=responses) mockModel = testing_utils.MockModel.create(responses=responses)
mockModel.model = 'gemini-1.5-pro' mockModel.model = 'gemini-1.5-pro'
# Calls the first time. # Calls the first time.
@ -45,12 +45,12 @@ def test_vertex_rag_retrieval_for_gemini_1_x():
) )
], ],
) )
runner = utils.InMemoryRunner(agent) runner = testing_utils.InMemoryRunner(agent)
events = runner.run('test1') events = runner.run('test1')
# Asserts the requests. # Asserts the requests.
assert len(mockModel.requests) == 1 assert len(mockModel.requests) == 1
assert utils.simplify_contents(mockModel.requests[0].contents) == [ assert testing_utils.simplify_contents(mockModel.requests[0].contents) == [
('user', 'test1'), ('user', 'test1'),
] ]
assert len(mockModel.requests[0].config.tools) == 1 assert len(mockModel.requests[0].config.tools) == 1
@ -65,7 +65,7 @@ def test_vertex_rag_retrieval_for_gemini_1_x_with_another_function_tool():
responses = [ responses = [
'response1', 'response1',
] ]
mockModel = utils.MockModel.create(responses=responses) mockModel = testing_utils.MockModel.create(responses=responses)
mockModel.model = 'gemini-1.5-pro' mockModel.model = 'gemini-1.5-pro'
# Calls the first time. # Calls the first time.
@ -83,12 +83,12 @@ def test_vertex_rag_retrieval_for_gemini_1_x_with_another_function_tool():
FunctionTool(func=noop_tool), FunctionTool(func=noop_tool),
], ],
) )
runner = utils.InMemoryRunner(agent) runner = testing_utils.InMemoryRunner(agent)
events = runner.run('test1') events = runner.run('test1')
# Asserts the requests. # Asserts the requests.
assert len(mockModel.requests) == 1 assert len(mockModel.requests) == 1
assert utils.simplify_contents(mockModel.requests[0].contents) == [ assert testing_utils.simplify_contents(mockModel.requests[0].contents) == [
('user', 'test1'), ('user', 'test1'),
] ]
assert len(mockModel.requests[0].config.tools[0].function_declarations) == 2 assert len(mockModel.requests[0].config.tools[0].function_declarations) == 2
@ -107,7 +107,7 @@ def test_vertex_rag_retrieval_for_gemini_2_x():
responses = [ responses = [
'response1', 'response1',
] ]
mockModel = utils.MockModel.create(responses=responses) mockModel = testing_utils.MockModel.create(responses=responses)
mockModel.model = 'gemini-2.0-flash' mockModel.model = 'gemini-2.0-flash'
# Calls the first time. # Calls the first time.
@ -124,12 +124,12 @@ def test_vertex_rag_retrieval_for_gemini_2_x():
) )
], ],
) )
runner = utils.InMemoryRunner(agent) runner = testing_utils.InMemoryRunner(agent)
events = runner.run('test1') events = runner.run('test1')
# Asserts the requests. # Asserts the requests.
assert len(mockModel.requests) == 1 assert len(mockModel.requests) == 1
assert utils.simplify_contents(mockModel.requests[0].contents) == [ assert testing_utils.simplify_contents(mockModel.requests[0].contents) == [
('user', 'test1'), ('user', 'test1'),
] ]
assert len(mockModel.requests[0].config.tools) == 1 assert len(mockModel.requests[0].config.tools) == 1

View File

@ -20,7 +20,7 @@ from pydantic import BaseModel
import pytest import pytest
from pytest import mark from pytest import mark
from .. import utils from .. import testing_utils
pytestmark = pytest.mark.skip( pytestmark = pytest.mark.skip(
reason='Skipping until tool.func evaluations are fixed (async)' reason='Skipping until tool.func evaluations are fixed (async)'
@ -50,7 +50,7 @@ def change_state_callback(callback_context: CallbackContext):
def test_no_schema(): def test_no_schema():
mock_model = utils.MockModel.create( mock_model = testing_utils.MockModel.create(
responses=[ responses=[
function_call_no_schema, function_call_no_schema,
'response1', 'response1',
@ -69,9 +69,9 @@ def test_no_schema():
tools=[AgentTool(agent=tool_agent)], tools=[AgentTool(agent=tool_agent)],
) )
runner = utils.InMemoryRunner(root_agent) runner = testing_utils.InMemoryRunner(root_agent)
assert utils.simplify_events(runner.run('test1')) == [ assert testing_utils.simplify_events(runner.run('test1')) == [
('root_agent', function_call_no_schema), ('root_agent', function_call_no_schema),
('root_agent', function_response_no_schema), ('root_agent', function_response_no_schema),
('root_agent', 'response2'), ('root_agent', 'response2'),
@ -81,7 +81,7 @@ def test_no_schema():
def test_update_state(): def test_update_state():
"""The agent tool can read and change parent state.""" """The agent tool can read and change parent state."""
mock_model = utils.MockModel.create( mock_model = testing_utils.MockModel.create(
responses=[ responses=[
function_call_no_schema, function_call_no_schema,
'{"custom_output": "response1"}', '{"custom_output": "response1"}',
@ -102,7 +102,7 @@ def test_update_state():
tools=[AgentTool(agent=tool_agent)], tools=[AgentTool(agent=tool_agent)],
) )
runner = utils.InMemoryRunner(root_agent) runner = testing_utils.InMemoryRunner(root_agent)
runner.session.state['state_1'] = 'state1_value' runner.session.state['state_1'] = 'state1_value'
runner.run('test1') runner.run('test1')
@ -128,7 +128,7 @@ def test_custom_schema():
class CustomOutput(BaseModel): class CustomOutput(BaseModel):
custom_output: str custom_output: str
mock_model = utils.MockModel.create( mock_model = testing_utils.MockModel.create(
responses=[ responses=[
function_call_custom, function_call_custom,
'{"custom_output": "response1"}', '{"custom_output": "response1"}',
@ -150,10 +150,10 @@ def test_custom_schema():
tools=[AgentTool(agent=tool_agent)], tools=[AgentTool(agent=tool_agent)],
) )
runner = utils.InMemoryRunner(root_agent) runner = testing_utils.InMemoryRunner(root_agent)
runner.session.state['state_1'] = 'state1_value' runner.session.state['state_1'] = 'state1_value'
assert utils.simplify_events(runner.run('test1')) == [ assert testing_utils.simplify_events(runner.run('test1')) == [
('root_agent', function_call_custom), ('root_agent', function_call_custom),
('root_agent', function_response_custom), ('root_agent', function_response_custom),
('root_agent', 'response2'), ('root_agent', 'response2'),

View File

@ -5,7 +5,7 @@ from google.adk.sessions import Session
from google.adk.utils import instructions_utils from google.adk.utils import instructions_utils
import pytest import pytest
from .. import utils from .. import testing_utils
class MockArtifactService: class MockArtifactService:
@ -32,7 +32,9 @@ async def _create_test_readonly_context(
name="agent", name="agent",
instruction="test", instruction="test",
) )
invocation_context = await utils.create_invocation_context(agent=agent) invocation_context = await testing_utils.create_invocation_context(
agent=agent
)
invocation_context.session = Session( invocation_context.session = Session(
state=state if state else {}, state=state if state else {},
app_name=app_name, app_name=app_name,