NeoAnd/adk-python
1
0
Fork 0
mirror of https://github.com/EvolutionAPI/adk-python.git synced 2025-07-13 07:04:51 -06:00
Code Issues Actions Packages Projects Releases Wiki Activity

Move public_utils to utils in tests

Browse Source 
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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tests/integration/test_callback.py
View File

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

6
tests/unittests/agents/test_base_agent.py
View File

@ -30,7 +30,7 @@ from google.genai import types
import pytest
import pytest_mock
from typing_extensions import override
from .. import utils
from .. import testing_utils
def _before_agent_callback_noop(callback_context: CallbackContext) -> None:
@ -398,7 +398,7 @@ async def test_before_agent_callbacks_chain(
request.function.__name__, agent
)
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)
for response in expected_responses
]
@ -459,7 +459,7 @@ async def test_after_agent_callbacks_chain(
request.function.__name__, agent
)
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)
for response in expected_responses
]

46
tests/unittests/agents/test_llm_agent_callbacks.py
View File

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

30
tests/unittests/agents/test_model_callback_chain.py
View File

@ -27,7 +27,7 @@ from google.genai import types
from pydantic import BaseModel
import pytest
from .. import utils
from .. import testing_utils
class CallbackType(Enum):
@ -42,7 +42,9 @@ async def mock_async_before_cb_side_effect(
):
if ret_value:
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
@ -54,7 +56,9 @@ def mock_sync_before_cb_side_effect(
):
if ret_value:
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
@ -66,7 +70,9 @@ async def mock_async_after_cb_side_effect(
):
if ret_value:
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
@ -78,7 +84,9 @@ def mock_sync_after_cb_side_effect(
):
if ret_value:
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
@ -129,7 +137,7 @@ async def test_before_model_callbacks_chain(
expected_calls: List[int],
):
responses = ["model_response"]
mock_model = utils.MockModel.create(responses=responses)
mock_model = testing_utils.MockModel.create(responses=responses)
mock_cbs = []
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],
)
runner = utils.TestInMemoryRunner(agent)
runner = testing_utils.TestInMemoryRunner(agent)
result = await runner.run_async_with_new_session("test")
assert utils.simplify_events(result) == [
assert testing_utils.simplify_events(result) == [
("root_agent", expected_response),
]
@ -191,7 +199,7 @@ async def test_after_model_callbacks_chain(
expected_calls: List[int],
):
responses = ["model_response"]
mock_model = utils.MockModel.create(responses=responses)
mock_model = testing_utils.MockModel.create(responses=responses)
mock_cbs = []
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],
)
runner = utils.TestInMemoryRunner(agent)
runner = testing_utils.TestInMemoryRunner(agent)
result = await runner.run_async_with_new_session("test")
assert utils.simplify_events(result) == [
assert testing_utils.simplify_events(result) == [
("root_agent", expected_response),
]

8
tests/unittests/flows/llm_flows/_test_examples.py
View File

@ -21,7 +21,7 @@ from google.adk.models.base_llm import LlmRequest
from google.genai import types
import pytest
from ... import utils
from ... import testing_utils
@pytest.mark.asyncio
@ -31,7 +31,7 @@ async def test_no_examples():
config=types.GenerateContentConfig(system_instruction=""),
)
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=""
)
@ -69,7 +69,7 @@ async def test_agent_examples():
name="agent",
examples=example_list,
)
invocation_context = await utils.create_invocation_context(
invocation_context = await testing_utils.create_invocation_context(
agent=agent, user_content="test"
)
@ -122,7 +122,7 @@ async def test_agent_base_example_provider():
name="agent",
examples=provider,
)
invocation_context = await utils.create_invocation_context(
invocation_context = await testing_utils.create_invocation_context(
agent=agent, user_content="test"
)

50
tests/unittests/flows/llm_flows/test_agent_transfer.py
View File

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

6
tests/unittests/flows/llm_flows/test_async_tool_callbacks.py
View File

@ -29,7 +29,7 @@ from google.adk.tools.tool_context import ToolContext
from google.genai import types
import pytest
from ... import utils
from ... import testing_utils
class CallbackType(Enum):
@ -73,7 +73,7 @@ async def invoke_tool_with_callbacks(
return {"initial": "response"}
tool = FunctionTool(simple_fn)
model = utils.MockModel.create(responses=[])
model = testing_utils.MockModel.create(responses=[])
agent = Agent(
name="agent",
model=model,
@ -81,7 +81,7 @@ async def invoke_tool_with_callbacks(
before_tool_callback=before_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=""
)
# Build function call event

58
tests/unittests/flows/llm_flows/test_functions_long_running.py
View File

@ -17,7 +17,7 @@ from google.adk.tools import ToolContext
from google.adk.tools.long_running_tool import LongRunningFunctionTool
from google.genai.types import Part
from ... import utils
from ... import testing_utils
def test_async_function():
@ -28,7 +28,7 @@ def test_async_function():
'response3',
'response4',
]
mockModel = utils.MockModel.create(responses=responses)
mockModel = testing_utils.MockModel.create(responses=responses)
function_called = 0
def increase_by_one(x: int, tool_context: ToolContext) -> int:
@ -43,14 +43,14 @@ def test_async_function():
model=mockModel,
tools=[LongRunningFunctionTool(func=increase_by_one)],
)
runner = utils.InMemoryRunner(agent)
runner = testing_utils.InMemoryRunner(agent)
events = runner.run('test1')
# Asserts the requests.
assert len(mockModel.requests) == 2
# 1 item: user content
assert mockModel.requests[0].contents == [
utils.UserContent('test1'),
testing_utils.UserContent('test1'),
]
increase_by_one_call = Part.from_function_call(
name='increase_by_one', args={'x': 1}
@ -59,7 +59,7 @@ def test_async_function():
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'),
('model', increase_by_one_call),
('user', pending_response),
@ -69,7 +69,7 @@ def test_async_function():
assert function_called == 1
# Asserts the responses.
assert utils.simplify_events(events) == [
assert testing_utils.simplify_events(events) == [
(
'root_agent',
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(
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.
assert len(mockModel.requests) == 3
assert utils.simplify_contents(mockModel.requests[2].contents) == [
assert testing_utils.simplify_contents(mockModel.requests[2].contents) == [
('user', 'test1'),
('model', increase_by_one_call),
('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.
result_response = Part.from_function_response(
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.
assert len(mockModel.requests) == 4
assert utils.simplify_contents(mockModel.requests[3].contents) == [
assert testing_utils.simplify_contents(mockModel.requests[3].contents) == [
('user', 'test1'),
('model', increase_by_one_call),
('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
# another summarization. Whether this is the right behavior is TBD.
another_result_response = Part.from_function_response(
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.
assert len(mockModel.requests) == 5
assert utils.simplify_contents(mockModel.requests[4].contents) == [
assert testing_utils.simplify_contents(mockModel.requests[4].contents) == [
('user', 'test1'),
('model', increase_by_one_call),
('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.
assert function_called == 1
@ -140,7 +140,7 @@ def test_async_function_with_none_response():
'response3',
'response4',
]
mockModel = utils.MockModel.create(responses=responses)
mockModel = testing_utils.MockModel.create(responses=responses)
function_called = 0
def increase_by_one(x: int, tool_context: ToolContext) -> int:
@ -154,20 +154,20 @@ def test_async_function_with_none_response():
model=mockModel,
tools=[LongRunningFunctionTool(func=increase_by_one)],
)
runner = utils.InMemoryRunner(agent)
runner = testing_utils.InMemoryRunner(agent)
events = runner.run('test1')
# Asserts the requests.
assert len(mockModel.requests) == 2
# 1 item: user content
assert mockModel.requests[0].contents == [
utils.UserContent('test1'),
testing_utils.UserContent('test1'),
]
increase_by_one_call = Part.from_function_call(
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'),
('model', increase_by_one_call),
(
@ -182,7 +182,7 @@ def test_async_function_with_none_response():
assert function_called == 1
# Asserts the responses.
assert utils.simplify_events(events) == [
assert testing_utils.simplify_events(events) == [
(
'root_agent',
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(
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.
assert len(mockModel.requests) == 3
assert utils.simplify_contents(mockModel.requests[2].contents) == [
assert testing_utils.simplify_contents(mockModel.requests[2].contents) == [
('user', 'test1'),
('model', increase_by_one_call),
('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.
result_response = Part.from_function_response(
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.
assert len(mockModel.requests) == 4
assert utils.simplify_contents(mockModel.requests[3].contents) == [
assert testing_utils.simplify_contents(mockModel.requests[3].contents) == [
('user', 'test1'),
('model', increase_by_one_call),
('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
# another summarization. Whether this is the right behavior is TBD.
another_result_response = Part.from_function_response(
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.
assert len(mockModel.requests) == 5
assert utils.simplify_contents(mockModel.requests[4].contents) == [
assert testing_utils.simplify_contents(mockModel.requests[4].contents) == [
('user', 'test1'),
('model', increase_by_one_call),
('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.
assert function_called == 1

10
tests/unittests/flows/llm_flows/test_functions_request_euc.py
View File

@ -28,7 +28,7 @@ from google.adk.tools import AuthToolArguments
from google.adk.tools import ToolContext
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:
@ -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]:
tool_context.request_credential(auth_config1)
@ -108,7 +108,7 @@ def test_function_request_euc():
model=mock_model,
tools=[call_external_api1, call_external_api2],
)
runner = utils.InMemoryRunner(agent)
runner = testing_utils.InMemoryRunner(agent)
events = runner.run('test')
assert events[0].content.parts[0].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
auth_config1 = AuthConfig(
@ -307,7 +307,7 @@ def test_function_get_auth_response():
model=mock_model,
tools=[call_external_api1, call_external_api2],
)
runner = utils.InMemoryRunner(agent)
runner = testing_utils.InMemoryRunner(agent)
runner.run('test')
request_euc_function_call_event = runner.session.events[-3]
function_response1 = types.FunctionResponse(

16
tests/unittests/flows/llm_flows/test_functions_sequential.py
View File

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

40
tests/unittests/flows/llm_flows/test_functions_simple.py
View File

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

10
tests/unittests/flows/llm_flows/test_identity.py
View File

@ -18,7 +18,7 @@ from google.adk.models import LlmRequest
from google.genai import types
import pytest
from ... import utils
from ... import testing_utils
@pytest.mark.asyncio
@ -28,7 +28,9 @@ async def test_no_description():
config=types.GenerateContentConfig(system_instruction=""),
)
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(
invocation_context,
@ -52,7 +54,9 @@ async def test_with_description():
name="agent",
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(
invocation_context,

30
tests/unittests/flows/llm_flows/test_instructions.py
View File

@ -20,7 +20,7 @@ from google.adk.sessions import Session
from google.genai import types
import pytest
from ... import utils
from ... import testing_utils
@pytest.mark.asyncio
@ -36,7 +36,9 @@ async def test_build_system_instruction():
{{customer_int }, { non-identifier-float}}, \
{'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(
app_name="test_app",
user_id="test_user",
@ -75,7 +77,9 @@ async def test_function_system_instruction():
name="agent",
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(
app_name="test_app",
user_id="test_user",
@ -117,7 +121,9 @@ async def test_async_function_system_instruction():
name="agent",
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(
app_name="test_app",
user_id="test_user",
@ -156,7 +162,9 @@ async def test_global_system_instruction():
model="gemini-1.5-flash",
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(
app_name="test_app",
user_id="test_user",
@ -198,7 +206,9 @@ async def test_function_global_system_instruction():
model="gemini-1.5-flash",
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(
app_name="test_app",
user_id="test_user",
@ -240,7 +250,9 @@ async def test_async_function_global_system_instruction():
model="gemini-1.5-flash",
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(
app_name="test_app",
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}."""
),
)
invocation_context = await utils.create_invocation_context(agent=agent)
invocation_context = await testing_utils.create_invocation_context(
agent=agent
)
invocation_context.session = Session(
app_name="test_app",
user_id="test_user",

36
tests/unittests/flows/llm_flows/test_model_callbacks.py
View File

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

8
tests/unittests/flows/llm_flows/test_other_configs.py
View File

@ -17,7 +17,7 @@ from google.adk.tools import ToolContext
from google.genai.types import Part
from pydantic import BaseModel
from ... import utils
from ... import testing_utils
def test_output_schema():
@ -27,7 +27,7 @@ def test_output_schema():
response = [
'response1',
]
mockModel = utils.MockModel.create(responses=response)
mockModel = testing_utils.MockModel.create(responses=response)
root_agent = Agent(
name='root_agent',
model=mockModel,
@ -36,9 +36,9 @@ def test_output_schema():
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'),
]
assert len(mockModel.requests) == 1

38
tests/unittests/flows/llm_flows/test_tool_callbacks.py
View File

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

6
tests/unittests/streaming/test_streaming.py
View File

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

0
tests/unittests/utils.py → tests/unittests/testing_utils.py
View File

20
tests/unittests/tools/retrieval/test_vertex_ai_rag_retrieval.py
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.genai import types
from ... import utils
from ... import testing_utils
def noop_tool(x: str) -> str:
@ -28,7 +28,7 @@ def test_vertex_rag_retrieval_for_gemini_1_x():
responses = [
'response1',
]
mockModel = utils.MockModel.create(responses=responses)
mockModel = testing_utils.MockModel.create(responses=responses)
mockModel.model = 'gemini-1.5-pro'
# 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')
# Asserts the requests.
assert len(mockModel.requests) == 1
assert utils.simplify_contents(mockModel.requests[0].contents) == [
assert testing_utils.simplify_contents(mockModel.requests[0].contents) == [
('user', 'test1'),
]
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 = [
'response1',
]
mockModel = utils.MockModel.create(responses=responses)
mockModel = testing_utils.MockModel.create(responses=responses)
mockModel.model = 'gemini-1.5-pro'
# 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),
],
)
runner = utils.InMemoryRunner(agent)
runner = testing_utils.InMemoryRunner(agent)
events = runner.run('test1')
# Asserts the requests.
assert len(mockModel.requests) == 1
assert utils.simplify_contents(mockModel.requests[0].contents) == [
assert testing_utils.simplify_contents(mockModel.requests[0].contents) == [
('user', 'test1'),
]
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 = [
'response1',
]
mockModel = utils.MockModel.create(responses=responses)
mockModel = testing_utils.MockModel.create(responses=responses)
mockModel.model = 'gemini-2.0-flash'
# 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')
# Asserts the requests.
assert len(mockModel.requests) == 1
assert utils.simplify_contents(mockModel.requests[0].contents) == [
assert testing_utils.simplify_contents(mockModel.requests[0].contents) == [
('user', 'test1'),
]
assert len(mockModel.requests[0].config.tools) == 1

18
tests/unittests/tools/test_agent_tool.py
View File

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

0
tests/unittests/public_utils/__init__.py → tests/unittests/utils/__init__.py
View File

6
tests/unittests/public_utils/test_instructions_utils.py → tests/unittests/utils/test_instructions_utils.py
View File

@ -5,7 +5,7 @@ from google.adk.sessions import Session
from google.adk.utils import instructions_utils
import pytest
from .. import utils
from .. import testing_utils
class MockArtifactService:
@ -32,7 +32,9 @@ async def _create_test_readonly_context(
name="agent",
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(
state=state if state else {},
app_name=app_name,