Move public_utils to utils in tests

Renamed conflicting utils.py as testing_utils.py

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

View File

@@ -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
]

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'),

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),
]