Skip to content

Commit b2eec54

Browse files
sararobcopybara-github
sararob
authored and
copybara-github
committed
chore: Add replay tests for private AE memories and memories_revisions modules
PiperOrigin-RevId: 885062735
1 parent 1b4c325 commit b2eec54

13 files changed

Lines changed: 434 additions & 0 deletions

‎tests/unit/vertexai/genai/replays/test_ae_memories_delete.py‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
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+
from tests.unit.vertexai.genai.replays import pytest_helper
18+
from vertexai._genai import types
19+
20+
21+
def test_delete_memory(client):
22+
ae_memory_operation = client.agent_engines.memories.delete(
23+
name="projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/memories/5605466683931099136",
24+
)
25+
assert isinstance(ae_memory_operation, types.DeleteAgentEngineMemoryOperation)
26+
27+
28+
pytestmark = pytest_helper.setup(
29+
file=__file__,
30+
globals_for_file=globals(),
31+
test_method="agent_engines.memories.delete",
32+
)

‎tests/unit/vertexai/genai/replays/test_ae_memories_get.py‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
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+
from tests.unit.vertexai.genai.replays import pytest_helper
18+
from vertexai._genai import types
19+
20+
21+
def test_get_memory(client):
22+
memory_name = "projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/memories/3858070028511346688"
23+
ae_memory = client.agent_engines.memories.get(name=memory_name)
24+
assert isinstance(ae_memory, types.Memory)
25+
assert ae_memory.name == memory_name
26+
27+
28+
pytestmark = pytest_helper.setup(
29+
file=__file__,
30+
globals_for_file=globals(),
31+
test_method="agent_engines.memories.get",
32+
)

‎tests/unit/vertexai/genai/replays/test_ae_memories_private_create.py‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
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+
from tests.unit.vertexai.genai.replays import pytest_helper
18+
from vertexai._genai import types
19+
20+
21+
def test_private_create_memory(client):
22+
ae_memory_operation = client.agent_engines.memories._create(
23+
name="projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584",
24+
fact="memory_fact",
25+
scope={"user_id": "123"},
26+
)
27+
assert isinstance(ae_memory_operation, types.AgentEngineMemoryOperation)
28+
29+
30+
pytestmark = pytest_helper.setup(
31+
file=__file__,
32+
globals_for_file=globals(),
33+
test_method="agent_engines.memories._create",
34+
)

‎tests/unit/vertexai/genai/replays/test_ae_memories_private_generate.py‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
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+
from tests.unit.vertexai.genai.replays import pytest_helper
18+
from vertexai._genai import types
19+
20+
21+
def test_private_generate_memory(client):
22+
ae_memory_operation = client.agent_engines.memories._generate(
23+
name="projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584",
24+
vertex_session_source=types.GenerateMemoriesRequestVertexSessionSource(
25+
session="projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/sessions/6922431337672474624"
26+
),
27+
)
28+
assert isinstance(ae_memory_operation, types.AgentEngineGenerateMemoriesOperation)
29+
30+
31+
pytestmark = pytest_helper.setup(
32+
file=__file__,
33+
globals_for_file=globals(),
34+
test_method="agent_engines.memories._generate",
35+
)

‎tests/unit/vertexai/genai/replays/test_ae_memories_private_get_generate_memories_operation.py‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
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+
from tests.unit.vertexai.genai.replays import pytest_helper
18+
from vertexai._genai import types
19+
20+
21+
def test_private_get_generate_memories_operation(client):
22+
memory_operation = client.agent_engines.memories._get_generate_memories_operation(
23+
operation_name="projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/operations/5669315676343369728"
24+
)
25+
assert isinstance(memory_operation, types.AgentEngineGenerateMemoriesOperation)
26+
27+
28+
pytestmark = pytest_helper.setup(
29+
file=__file__,
30+
globals_for_file=globals(),
31+
test_method="agent_engines.memories._get_generate_memories_operation",
32+
)

‎tests/unit/vertexai/genai/replays/test_ae_memories_private_get_memory_operation.py‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
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+
from tests.unit.vertexai.genai.replays import pytest_helper
18+
from vertexai._genai import types
19+
20+
21+
def test_private_get_memory_operation(client):
22+
memory_operation = client.agent_engines.memories._get_memory_operation(
23+
operation_name="projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/memories/3858070028511346688/operations/1044963283964002304"
24+
)
25+
assert isinstance(memory_operation, types.AgentEngineMemoryOperation)
26+
27+
28+
pytestmark = pytest_helper.setup(
29+
file=__file__,
30+
globals_for_file=globals(),
31+
test_method="agent_engines.memories._get_memory_operation",
32+
)

‎tests/unit/vertexai/genai/replays/test_ae_memories_private_list.py‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
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+
from tests.unit.vertexai.genai.replays import pytest_helper
18+
from vertexai._genai import types
19+
20+
21+
def test_private_list_memory(client):
22+
ae_name = "projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584"
23+
memory_list = client.agent_engines.memories._list(name=ae_name)
24+
assert isinstance(memory_list, types.ListReasoningEnginesMemoriesResponse)
25+
assert isinstance(memory_list.memories[0], types.Memory)
26+
27+
28+
pytestmark = pytest_helper.setup(
29+
file=__file__,
30+
globals_for_file=globals(),
31+
test_method="agent_engines.memories._list",
32+
)

‎tests/unit/vertexai/genai/replays/test_ae_memories_private_purge.py‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
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+
from tests.unit.vertexai.genai.replays import pytest_helper
18+
from vertexai._genai import types
19+
20+
21+
def test_private_purge(client):
22+
ae_memory_purge_operation = client.agent_engines.memories._purge(
23+
name="projects/964831358985/locations/us-central1/reasoningEngines/6086402690647064576",
24+
filter="scope.user_id=123",
25+
)
26+
assert isinstance(
27+
ae_memory_purge_operation, types.AgentEnginePurgeMemoriesOperation
28+
)
29+
30+
31+
pytestmark = pytest_helper.setup(
32+
file=__file__,
33+
globals_for_file=globals(),
34+
test_method="agent_engines.memories._purge",
35+
)

‎tests/unit/vertexai/genai/replays/test_ae_memories_private_retrieve.py‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
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+
from tests.unit.vertexai.genai.replays import pytest_helper
18+
from vertexai._genai import types
19+
20+
21+
def test_private_retrieve(client):
22+
ae_name = "projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584"
23+
retrieved_memories = client.agent_engines.memories._retrieve(
24+
name=ae_name,
25+
scope={"user_id": "123"},
26+
)
27+
assert isinstance(retrieved_memories, types.RetrieveMemoriesResponse)
28+
assert isinstance(
29+
retrieved_memories.retrieved_memories[0],
30+
types.RetrieveMemoriesResponseRetrievedMemory,
31+
)
32+
33+
34+
pytestmark = pytest_helper.setup(
35+
file=__file__,
36+
globals_for_file=globals(),
37+
test_method="agent_engines.memories._retrieve",
38+
)

‎tests/unit/vertexai/genai/replays/test_ae_memories_private_rollback.py‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
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+
from tests.unit.vertexai.genai.replays import pytest_helper
18+
from vertexai._genai import types
19+
20+
21+
def test_private_rollback(client):
22+
rollback_operation = client.agent_engines.memories._rollback(
23+
name="projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/memories/3858070028511346688",
24+
target_revision_id="3001207491565453312",
25+
)
26+
assert isinstance(rollback_operation, types.AgentEngineRollbackMemoryOperation)
27+
28+
29+
pytestmark = pytest_helper.setup(
30+
file=__file__,
31+
globals_for_file=globals(),
32+
test_method="agent_engines.memories._rollback",
33+
)

0 commit comments

Comments
 (0)