Skip to content

Commit 95ae4c8

Browse files
sararobcopybara-github
authored andcommitted
chore: GenAI SDK client - Add more replay tests for Agent Engines
PiperOrigin-RevId: 785892911
1 parent 82dd075 commit 95ae4c8

File tree

3 files changed

+90
-0
lines changed

3 files changed

+90
-0
lines changed
Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
#
15+
# pylint: disable=protected-access,bad-continuation,missing-function-docstring
16+
17+
import pytest
18+
19+
20+
from tests.unit.vertexai.genai.replays import pytest_helper
21+
from vertexai._genai import types
22+
23+
24+
def test_agent_engine_delete(client):
25+
agent_engine = client.agent_engines.create()
26+
operation = client.agent_engines.delete(name=agent_engine.api_resource.name)
27+
assert isinstance(operation, types.DeleteAgentEngineOperation)
28+
29+
30+
pytestmark = pytest_helper.setup(
31+
file=__file__,
32+
globals_for_file=globals(),
33+
test_method="agent_engines.delete",
34+
)
35+
36+
37+
pytest_plugins = ("pytest_asyncio",)
38+
39+
40+
@pytest.mark.asyncio
41+
async def test_agent_engine_delete_async(client):
42+
# TODO(b/431785750): use async methods for create() when available
43+
agent_engine = client.agent_engines.create()
44+
operation = await client.aio.agent_engines.delete(
45+
name=agent_engine.api_resource.name
46+
)
47+
assert isinstance(operation, types.DeleteAgentEngineOperation)

tests/unit/vertexai/genai/replays/test_delete_agent_engine_memory.py

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,9 @@
1414
#
1515
# pylint: disable=protected-access,bad-continuation,missing-function-docstring
1616

17+
import pytest
18+
19+
1720
from tests.unit.vertexai.genai.replays import pytest_helper
1821
from vertexai._genai import types
1922

@@ -36,3 +39,21 @@ def test_delete_memory(client):
3639
globals_for_file=globals(),
3740
test_method="agent_engines.delete_memory",
3841
)
42+
43+
44+
pytest_plugins = ("pytest_asyncio",)
45+
46+
47+
@pytest.mark.asyncio
48+
async def test_delete_memory_async(client):
49+
# TODO(b/431785750): use async methods for create() and create_memory() when available
50+
agent_engine = client.agent_engines.create()
51+
operation = client.agent_engines.create_memory(
52+
name=agent_engine.api_resource.name,
53+
fact="memory_fact",
54+
scope={"user_id": "123"},
55+
)
56+
memory = operation.response
57+
operation = await client.aio.agent_engines.delete_memory(name=memory.name)
58+
assert isinstance(operation, types.DeleteAgentEngineMemoryOperation)
59+
assert operation.name.startswith(memory.name + "/operations/")

tests/unit/vertexai/genai/replays/test_get_agent_engine_memory.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,8 @@
1414
#
1515
# pylint: disable=protected-access,bad-continuation,missing-function-docstring
1616

17+
import pytest
18+
1719
from tests.unit.vertexai.genai.replays import pytest_helper
1820
from vertexai._genai import types
1921

@@ -38,3 +40,23 @@ def test_get_memory(client):
3840
globals_for_file=globals(),
3941
test_method="agent_engines.get_memory",
4042
)
43+
44+
45+
pytest_plugins = ("pytest_asyncio",)
46+
47+
48+
@pytest.mark.asyncio
49+
async def test_get_memory_async(client):
50+
# TODO(b/431785750): use async methods for create() and create_memory() when available
51+
agent_engine = client.agent_engines.create()
52+
operation = client.agent_engines.create_memory(
53+
name=agent_engine.api_resource.name,
54+
fact="memory_fact",
55+
scope={"user_id": "123"},
56+
)
57+
assert isinstance(operation, types.AgentEngineMemoryOperation)
58+
memory = await client.aio.agent_engines.get_memory(
59+
name=operation.response.name,
60+
)
61+
assert isinstance(memory, types.Memory)
62+
assert memory.name == operation.response.name

0 commit comments

Comments
 (0)