Skip to content

Commit 458fa3f

Browse files
yeesiancopybara-github
authored andcommitted
chore: add deprecated methods to AdkApp
PiperOrigin-RevId: 825034004
1 parent 8075e60 commit 458fa3f

File tree

2 files changed

+354
-0
lines changed

2 files changed

+354
-0
lines changed

tests/unit/vertex_adk/test_agent_engine_templates_adk.py

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -264,6 +264,39 @@ def test_register_operations(self):
264264
for operation in operations:
265265
assert operation in dir(app)
266266

267+
def test_stream_query(self):
268+
app = agent_engines.AdkApp(agent=_TEST_AGENT)
269+
assert app._tmpl_attrs.get("runner") is None
270+
app.set_up()
271+
app._tmpl_attrs["runner"] = _MockRunner()
272+
events = list(
273+
app.stream_query(
274+
user_id=_TEST_USER_ID,
275+
message="test message",
276+
)
277+
)
278+
assert len(events) == 1
279+
280+
def test_stream_query_with_content(self):
281+
app = agent_engines.AdkApp(agent=_TEST_AGENT)
282+
assert app._tmpl_attrs.get("runner") is None
283+
app.set_up()
284+
app._tmpl_attrs["runner"] = _MockRunner()
285+
events = list(
286+
app.stream_query(
287+
user_id=_TEST_USER_ID,
288+
message=types.Content(
289+
role="user",
290+
parts=[
291+
types.Part(
292+
text="test message with content",
293+
)
294+
],
295+
).model_dump(),
296+
)
297+
)
298+
assert len(events) == 1
299+
267300
@pytest.mark.asyncio
268301
async def test_async_stream_query(self):
269302
app = agent_engines.AdkApp(agent=_TEST_AGENT)
@@ -385,6 +418,51 @@ async def test_async_delete_session(self):
385418
response0 = await app.async_list_sessions(user_id=_TEST_USER_ID)
386419
assert not response0.sessions
387420

421+
def test_create_session(self):
422+
app = agent_engines.AdkApp(agent=_TEST_AGENT)
423+
session1 = app.create_session(user_id=_TEST_USER_ID)
424+
assert session1.user_id == _TEST_USER_ID
425+
session2 = app.create_session(
426+
user_id=_TEST_USER_ID, session_id="test_session_id"
427+
)
428+
assert session2.user_id == _TEST_USER_ID
429+
assert session2.id == "test_session_id"
430+
431+
def test_get_session(self):
432+
app = agent_engines.AdkApp(agent=_TEST_AGENT)
433+
session1 = app.create_session(user_id=_TEST_USER_ID)
434+
session2 = app.get_session(
435+
user_id=_TEST_USER_ID,
436+
session_id=session1.id,
437+
)
438+
assert session2.user_id == _TEST_USER_ID
439+
assert session1.id == session2.id
440+
441+
def test_list_sessions(self):
442+
app = agent_engines.AdkApp(agent=_TEST_AGENT)
443+
response0 = app.list_sessions(user_id=_TEST_USER_ID)
444+
assert not response0.sessions
445+
session = app.create_session(user_id=_TEST_USER_ID)
446+
response1 = app.list_sessions(user_id=_TEST_USER_ID)
447+
assert len(response1.sessions) == 1
448+
assert response1.sessions[0].id == session.id
449+
session2 = app.create_session(user_id=_TEST_USER_ID)
450+
response2 = app.list_sessions(user_id=_TEST_USER_ID)
451+
assert len(response2.sessions) == 2
452+
assert response2.sessions[0].id == session.id
453+
assert response2.sessions[1].id == session2.id
454+
455+
def test_delete_session(self):
456+
app = agent_engines.AdkApp(agent=_TEST_AGENT)
457+
response = app.delete_session(user_id=_TEST_USER_ID, session_id="")
458+
assert not response
459+
session = app.create_session(user_id=_TEST_USER_ID)
460+
response1 = app.list_sessions(user_id=_TEST_USER_ID)
461+
assert len(response1.sessions) == 1
462+
app.delete_session(user_id=_TEST_USER_ID, session_id=session.id)
463+
response0 = app.list_sessions(user_id=_TEST_USER_ID)
464+
assert not response0.sessions
465+
388466
@pytest.mark.asyncio
389467
async def test_async_add_session_to_memory_dict(self):
390468
app = agent_engines.AdkApp(agent=_TEST_AGENT)

0 commit comments

Comments
 (0)