diff --git a/tests/integration/test_callback.py b/tests/integration/test_callback.py index 4d9f1d3..b211003 100644 --- a/tests/integration/test_callback.py +++ b/tests/integration/test_callback.py @@ -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 diff --git a/tests/unittests/agents/test_base_agent.py b/tests/unittests/agents/test_base_agent.py index 3378143..95048ea 100644 --- a/tests/unittests/agents/test_base_agent.py +++ b/tests/unittests/agents/test_base_agent.py @@ -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 ] diff --git a/tests/unittests/agents/test_llm_agent_callbacks.py b/tests/unittests/agents/test_llm_agent_callbacks.py index 99a606e..21ef8a9 100644 --- a/tests/unittests/agents/test_llm_agent_callbacks.py +++ b/tests/unittests/agents/test_llm_agent_callbacks.py @@ -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'), diff --git a/tests/unittests/agents/test_model_callback_chain.py b/tests/unittests/agents/test_model_callback_chain.py index 3457e1d..e0bf037 100644 --- a/tests/unittests/agents/test_model_callback_chain.py +++ b/tests/unittests/agents/test_model_callback_chain.py @@ -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), ] diff --git a/tests/unittests/flows/llm_flows/_test_examples.py b/tests/unittests/flows/llm_flows/_test_examples.py index 29eb718..c26116f 100644 --- a/tests/unittests/flows/llm_flows/_test_examples.py +++ b/tests/unittests/flows/llm_flows/_test_examples.py @@ -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" ) diff --git a/tests/unittests/flows/llm_flows/test_agent_transfer.py b/tests/unittests/flows/llm_flows/test_agent_transfer.py index f236077..6e2bfe7 100644 --- a/tests/unittests/flows/llm_flows/test_agent_transfer.py +++ b/tests/unittests/flows/llm_flows/test_agent_transfer.py @@ -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'), ] diff --git a/tests/unittests/flows/llm_flows/test_async_tool_callbacks.py b/tests/unittests/flows/llm_flows/test_async_tool_callbacks.py index 051bdde..35f3a81 100644 --- a/tests/unittests/flows/llm_flows/test_async_tool_callbacks.py +++ b/tests/unittests/flows/llm_flows/test_async_tool_callbacks.py @@ -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 diff --git a/tests/unittests/flows/llm_flows/test_functions_long_running.py b/tests/unittests/flows/llm_flows/test_functions_long_running.py index a547517..e173c87 100644 --- a/tests/unittests/flows/llm_flows/test_functions_long_running.py +++ b/tests/unittests/flows/llm_flows/test_functions_long_running.py @@ -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 diff --git a/tests/unittests/flows/llm_flows/test_functions_request_euc.py b/tests/unittests/flows/llm_flows/test_functions_request_euc.py index 6dcb6f9..3f5abca 100644 --- a/tests/unittests/flows/llm_flows/test_functions_request_euc.py +++ b/tests/unittests/flows/llm_flows/test_functions_request_euc.py @@ -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( diff --git a/tests/unittests/flows/llm_flows/test_functions_sequential.py b/tests/unittests/flows/llm_flows/test_functions_sequential.py index 02c2d41..0a21b8d 100644 --- a/tests/unittests/flows/llm_flows/test_functions_sequential.py +++ b/tests/unittests/flows/llm_flows/test_functions_sequential.py @@ -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})), diff --git a/tests/unittests/flows/llm_flows/test_functions_simple.py b/tests/unittests/flows/llm_flows/test_functions_simple.py index 0e9e43f..7ddb124 100644 --- a/tests/unittests/flows/llm_flows/test_functions_simple.py +++ b/tests/unittests/flows/llm_flows/test_functions_simple.py @@ -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: diff --git a/tests/unittests/flows/llm_flows/test_identity.py b/tests/unittests/flows/llm_flows/test_identity.py index 564400c..336da64 100644 --- a/tests/unittests/flows/llm_flows/test_identity.py +++ b/tests/unittests/flows/llm_flows/test_identity.py @@ -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, diff --git a/tests/unittests/flows/llm_flows/test_instructions.py b/tests/unittests/flows/llm_flows/test_instructions.py index d6e2234..8d95dc2 100644 --- a/tests/unittests/flows/llm_flows/test_instructions.py +++ b/tests/unittests/flows/llm_flows/test_instructions.py @@ -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", diff --git a/tests/unittests/flows/llm_flows/test_model_callbacks.py b/tests/unittests/flows/llm_flows/test_model_callbacks.py index dd2d3cf..154ee80 100644 --- a/tests/unittests/flows/llm_flows/test_model_callbacks.py +++ b/tests/unittests/flows/llm_flows/test_model_callbacks.py @@ -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')] diff --git a/tests/unittests/flows/llm_flows/test_other_configs.py b/tests/unittests/flows/llm_flows/test_other_configs.py index 63ba950..005bd1b 100644 --- a/tests/unittests/flows/llm_flows/test_other_configs.py +++ b/tests/unittests/flows/llm_flows/test_other_configs.py @@ -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 diff --git a/tests/unittests/flows/llm_flows/test_tool_callbacks.py b/tests/unittests/flows/llm_flows/test_tool_callbacks.py index 5383f41..1f26b18 100644 --- a/tests/unittests/flows/llm_flows/test_tool_callbacks.py +++ b/tests/unittests/flows/llm_flows/test_tool_callbacks.py @@ -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( diff --git a/tests/unittests/streaming/test_streaming.py b/tests/unittests/streaming/test_streaming.py index a20da0f..c1e1eaa 100644 --- a/tests/unittests/streaming/test_streaming.py +++ b/tests/unittests/streaming/test_streaming.py @@ -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() diff --git a/tests/unittests/utils.py b/tests/unittests/testing_utils.py similarity index 100% rename from tests/unittests/utils.py rename to tests/unittests/testing_utils.py diff --git a/tests/unittests/tools/retrieval/test_vertex_ai_rag_retrieval.py b/tests/unittests/tools/retrieval/test_vertex_ai_rag_retrieval.py index f8d122c..b55cfe1 100644 --- a/tests/unittests/tools/retrieval/test_vertex_ai_rag_retrieval.py +++ b/tests/unittests/tools/retrieval/test_vertex_ai_rag_retrieval.py @@ -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 diff --git a/tests/unittests/tools/test_agent_tool.py b/tests/unittests/tools/test_agent_tool.py index dc8cdeb..36a8158 100644 --- a/tests/unittests/tools/test_agent_tool.py +++ b/tests/unittests/tools/test_agent_tool.py @@ -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'), diff --git a/tests/unittests/public_utils/__init__.py b/tests/unittests/utils/__init__.py similarity index 100% rename from tests/unittests/public_utils/__init__.py rename to tests/unittests/utils/__init__.py diff --git a/tests/unittests/public_utils/test_instructions_utils.py b/tests/unittests/utils/test_instructions_utils.py similarity index 98% rename from tests/unittests/public_utils/test_instructions_utils.py rename to tests/unittests/utils/test_instructions_utils.py index 8fc7647..35e5195 100644 --- a/tests/unittests/public_utils/test_instructions_utils.py +++ b/tests/unittests/utils/test_instructions_utils.py @@ -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,