8000 Move public_utils to utils in tests · rondweb/adk-python@41b33d4 · GitHub
[go: up one dir, main page]
Skip to content

Commit 41b33d4

Browse files
selcukguncopybara-github
selcukgun
authored and
copybara-github
committed
Move public_utils to utils in tests
Renamed conflicting utils.py as testing_utils.py PiperOrigin-RevId: 761715808
1 parent 09cb128 commit 41b33d4

22 files changed

+236
-208
lines changed

‎tests/integration/test_callback.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414

1515
from pytest import mark
1616

17-
from ..unittests.utils import simplify_events
17+
from ..unittests.testing_utils import simplify_events
1818
from .fixture import callback_agent
1919
from .utils import assert_agent_says
2020
from .utils import TestRunner

‎tests/unittests/agents/test_base_agent.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@
3030
import pytest
3131
import pytest_mock
3232
from typing_extensions import override
33-
from .. import utils
33+
from .. import testing_utils
3434

3535

3636
def _before_agent_callback_noop(callback_context: CallbackContext) -> None:
@@ -398,7 +398,7 @@ async def test_before_agent_callbacks_chain(
398398
request.function.__name__, agent
399399
)
400400
result = [e async for e in agent.run_async(parent_ctx)]
401-
assert utils.simplify_events(result) == [
401+
assert testing_utils.simplify_events(result) == [
402402
(f'{request.function.__name__}_test_agent', response)
403403
for response in expected_responses
404404
]
@@ -459,7 +459,7 @@ async def test_after_agent_callbacks_chain(
459459
request.function.__name__, agent
460460
)
461461
result = [e async for e in agent.run_async(parent_ctx)]
462-
assert utils.simplify_events(result) == [
462+
assert testing_utils.simplify_events(result) == [
463463
(f'{request.function.__name__}_test_agent', response)
464464
for response in expected_responses
465465
]

‎tests/unittests/agents/test_llm_agent_callbacks.py

Lines changed: 23 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@
2323
from pydantic import BaseModel
2424
import pytest
2525

26-
from .. import utils
26+
from .. import testing_utils
2727

2828

2929
class MockBeforeModelCallback(BaseModel):
@@ -35,7 +35,7 @@ def __call__(
3535
llm_request: LlmRequest,
3636
) -> LlmResponse:
3737
return LlmResponse(
38-
content=utils.ModelContent(
38+
content=testing_utils.ModelContent(
3939
[types.Part.from_text(text=self.mock_response)]
4040
)
4141
)
@@ -50,7 +50,7 @@ def __call__(
5050
llm_response: LlmResponse,
5151
) -> LlmResponse:
5252
return LlmResponse(
53-
content=utils.ModelContent(
53+
content=testing_utils.ModelContent(
5454
[types.Part.from_text(text=self.mock_response)]
5555
)
5656
)
@@ -65,7 +65,7 @@ async def __call__(
6565
llm_request: LlmRequest,
6666
) -> LlmResponse:
6767
return LlmResponse(
68-
content=utils.ModelContent(
68+
content=testing_utils.ModelContent(
6969
[types.Part.from_text(text=self.mock_response)]
7070
)
7171
)
@@ -80,7 +80,7 @@ async def __call__(
8080
llm_response: LlmResponse,
8181
) -> LlmResponse:
8282
return LlmResponse(
83-
content=utils.ModelContent(
83+
content=testing_utils.ModelContent(
8484
[types.Part.from_text(text=self.mock_response)]
8585
)
8686
)
@@ -97,7 +97,7 @@ async def async_noop_callback(**kwargs) -> Optional[LlmResponse]:
9797
@pytest.mark.asyncio
9898
async def test_before_model_callback():
9999
responses = ['model_response']
100-
mock_model = utils.MockModel.create(responses=responses)
100+
mock_model = testing_utils.MockModel.create(responses=responses)
101101
agent = Agent(
102102
name='root_agent',
103103
model=mock_model,
@@ -106,8 +106,8 @@ async def test_before_model_callback():
106106
),
107107
)
108108

109-
runner = utils.TestInMemoryRunner(agent)
110-
assert utils.simplify_events(
109+
runner = testing_utils.TestInMemoryRunner(agent)
110+
assert testing_utils.simplify_events(
111111
await runner.run_async_with_new_session('test')
112112
) == [
113113
('root_agent', 'before_model_callback'),
@@ -117,15 +117,15 @@ async def test_before_model_callback():
117117
@pytest.mark.asyncio
118118
async def test_before_model_callback_noop():
119119
responses = ['model_response']
120-
mock_model = utils.MockModel.create(responses=responses)
120+
mock_model = testing_utils.MockModel.create(responses=responses)
121121
agent = Agent(
122122
name='root_agent',
123123
model=mock_model,
124124
before_model_callback=noop_callback,
125125
)
126126

127-
runner = utils.TestInMemoryRunner(agent)
128-
assert utils.simplify_events(
127+
runner = testing_utils.TestInMemoryRunner(agent)
128+
assert testing_utils.simplify_events(
129129
await runner.run_async_with_new_session('test')
130130
) == [
131131
('root_agent', 'model_response'),
@@ -135,7 +135,7 @@ async def test_before_model_callback_noop():
135135
@pytest.mark.asyncio
136136
async def test_after_model_callback():
137137
responses = ['model_response']
138-
mock_model = utils.MockModel.create(responses=responses)
138+
mock_model = testing_utils.MockModel.create(responses=responses)
139139
agent = Agent(
140140
name='root_agent',
141141
model=mock_model,
@@ -144,8 +144,8 @@ async def test_after_model_callback():
144144
),
145145
)
146146

147-
runner = utils.TestInMemoryRunner(agent)
148-
assert utils.simplify_events(
147+
runner = testing_utils.TestInMemoryRunner(agent)
148+
assert testing_utils.simplify_events(
149149
await runner.run_async_with_new_session('test')
150150
) == [
151151
('root_agent', 'after_model_callback'),
@@ -155,7 +155,7 @@ async def test_after_model_callback():
155155
@pytest.mark.asyncio
156156
async def test_async_before_model_callback():
157157
responses = ['model_response']
158-
mock_model = utils.MockModel.create(responses=responses)
158+
mock_model = testing_utils.MockModel.create(responses=responses)
159159
agent = Agent(
160160
name='root_agent',
161161
model=mock_model,
@@ -164,8 +164,8 @@ async def test_async_before_model_callback():
164164
),
165165
)
166166

167-
runner = utils.TestInMemoryRunner(agent)
168-
assert utils.simplify_events(
167+
runner = testing_utils.TestInMemoryRunner(agent)
168+
assert testing_utils.simplify_events(
169169
await runner.run_async_with_new_session('test')
170170
) == [
171171
('root_agent', 'async_before_model_callback'),
@@ -175,15 +175,15 @@ async def test_async_before_model_callback():
175175
@pytest.mark.asyncio
176176
async def test_async_before_model_callback_noop():
177177
responses = ['model_response']
178-
mock_model = utils.MockModel.create(responses=responses)
178+
mock_model = testing_utils.MockModel.create(responses=responses)
179179
agent = Agent(
180180
name='root_agent',
181181
model=mock_model,
182182
before_model_callback=async_noop_callback,
183183
)
184184

185-
runner = utils.TestInMemoryRunner(agent)
186-
assert utils.simplify_events(
185+
runner = testing_utils.TestInMemoryRunner(agent)
186+
assert testing_utils.simplify_events(
187187
await runner.run_async_with_new_session('test')
188188
) == [
189189
('root_agent', 'model_response'),
@@ -193,7 +193,7 @@ async def test_async_before_model_callback_noop():
193193
@pytest.mark.asyncio
194194
async def test_async_after_model_callback():
195195
responses = ['model_response']
196-
mock_model = utils.MockModel.create(responses=responses)
196+
mock_model = testing_utils.MockModel.create(responses=responses)
197197
agent = Agent(
198198
name='root_agent',
199199
model=mock_model,
@@ -202,8 +202,8 @@ async def test_async_after_model_callback():
202202
),
203203
)
204204

205-
runner = utils.TestInMemoryRunner(agent)
206-
assert utils.simplify_events(
205+
runner = testing_utils.TestInMemoryRunner(agent)
206+
assert testing_utils.simplify_events(
207207
await runner.run_async_with_new_session('test')
208208
) == [
209209
('root_agent', 'async_after_model_callback'),

‎tests/unittests/agents/test_model_callback_chain.py

Lines changed: 19 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@
2727
from pydantic import BaseModel
2828
import pytest
2929

30-
from .. import utils
30+
from .. import testing_utils
3131

3232

3333
class CallbackType(Enum):
@@ -42,7 +42,9 @@ async def mock_async_before_cb_side_effect(
4242
):
4343
if ret_value:
4444
return LlmResponse(
45-
content=utils.ModelContent([types.Part.from_text(text=ret_value)])
45+
content=testing_utils.ModelContent(
46+
[types.Part.from_text(text=ret_value)]
47+
)
4648
)
4749
return None
4850

@@ -54,7 +56,9 @@ def mock_sync_before_cb_side_effect(
5456
):
5557
if ret_value:
5658
return LlmResponse(
57-
content=utils.ModelContent([types.Part.from_text(text=ret_value)])
59+
content=testing_utils.ModelContent(
60+
[types.Part.from_text(text=ret_value)]
61+
)
5862
)
5963
return None
6064

@@ -66,7 +70,9 @@ async def mock_async_after_cb_side_effect(
6670
):
6771
if ret_value:
6872
return LlmResponse(
69-
content=utils.ModelContent([types.Part.from_text(text=ret_value)])
73+
content=testing_utils.ModelContent(
74+
[types.Part.from_text(text=ret_value)]
75+
)
7076
)
7177
return None
7278

@@ -78,7 +84,9 @@ def mock_sync_after_cb_side_effect(
7884
):
7985
if ret_value:
8086
return LlmResponse(
81-
content=utils.ModelContent([types.Part.from_text(text=ret_value)])
87+
content=testing_utils.ModelContent(
88+
[types.Part.from_text(text=ret_value)]
89+
)
8290
)
8391
return None
8492

@@ -129,7 +137,7 @@ async def test_before_model_callbacks_chain(
129137
expected_calls: List[int],
130138
):
131139
responses = ["model_response"]
132-
mock_model = utils.MockModel.create(responses=responses)
140+
mock_model = testing_utils.MockModel.create(responses=responses)
133141

134142
mock_cbs = []
135143
for response, callback_type in callbacks:
@@ -154,9 +162,9 @@ async def test_before_model_callbacks_chain(
154162
before_model_callback=[mock_cb for mock_cb in mock_cbs],
155163
)
156164

157-
runner = utils.TestInMemoryRunner(agent)
165+
runner = testing_utils.TestInMemoryRunner(agent)
158166
result = await runner.run_async_with_new_session("test")
159-
assert utils.simplify_events(result) == [
167+
assert testing_utils.simplify_events(result) == [
160168
("root_agent", expected_response),
161169
]
162170

@@ -191,7 +199,7 @@ async def test_after_model_callbacks_chain(
191199
expected_calls: List[int],
192200
):
193201
responses = ["model_response"]
194-
mock_model = utils.MockModel.create(responses=responses)
202+
mock_model = testing_utils.MockModel.create(responses=responses)
195203

196204
mock_cbs = []
197205
for response, callback_type in callbacks:
@@ -216,9 +224,9 @@ async def test_after_model_callbacks_chain(
216224
after_model_callback=[mock_cb for mock_cb in mock_cbs],
217225
)
218226

219-
runner = utils.TestInMemoryRunner(agent)
227+
runner = testing_utils.TestInMemoryRunner(agent)
220228
result = await runner.run_async_with_new_session("test")
221-
assert utils.simplify_events(result) == [
229+
assert testing_utils.simplify_events(result) == [
222230
("root_agent", expected_response),
223231
]
224232

‎tests/unittests/flows/llm_flows/_test_examples.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@
2121
from google.genai import types
2222
import pytest
2323

24-
from ... import utils
24+
from ... import testing_utils
2525

2626

2727
@pytest.mark.asyncio
@@ -31,7 +31,7 @@ async def test_no_examples():
3131
config=types.GenerateContentConfig(system_instruction=""),
3232
)
3333
agent = Agent(model="gemini-1.5-flash", name="agent", examples=[])
34-
invocation_context = await utils.create_invocation_context(
34+
invocation_context = await testing_utils.create_invocation_context(
3535
agent=agent, user_content=""
3636
)
3737

@@ -69,7 +69,7 @@ async def test_agent_examples():
6969
name="agent",
7070
examples=example_list,
7171
)
72-
invocation_context = await utils.create_invocation_context(
72+
invocation_context = await testing_utils.create_invocation_context(
7373
agent=agent, user_content="test"
7474
)
7575

@@ -122,7 +122,7 @@ def get_examples(self, query: str) -> list[Example]:
122122
name="agent",
123123
examples=provider,
124124
)
125-
invocation_context = await utils.create_invocation_context(
125+
invocation_context = await testing_utils.create_invocation_context(
126126
agent=agent, user_content="test"
127127
)
128128

0 commit comments

Comments
 (0)
0