From c13769e5472fb4f5d83163055a7d21ac19bf6e74 Mon Sep 17 00:00:00 2001 From: Sara Robinson Date: Tue, 17 Mar 2026 09:30:15 -0700 Subject: [PATCH] chore: Add replay tests for private AE memories and memories_revisions modules PiperOrigin-RevId: 885062735 --- .../genai/replays/test_ae_memories_delete.py | 32 ++++++++++++++++ .../genai/replays/test_ae_memories_get.py | 32 ++++++++++++++++ .../test_ae_memories_private_create.py | 34 +++++++++++++++++ .../test_ae_memories_private_generate.py | 35 +++++++++++++++++ ...private_get_generate_memories_operation.py | 32 ++++++++++++++++ ...e_memories_private_get_memory_operation.py | 32 ++++++++++++++++ .../replays/test_ae_memories_private_list.py | 32 ++++++++++++++++ .../replays/test_ae_memories_private_purge.py | 35 +++++++++++++++++ .../test_ae_memories_private_retrieve.py | 38 +++++++++++++++++++ .../test_ae_memories_private_rollback.py | 33 ++++++++++++++++ .../test_ae_memories_private_update.py | 35 +++++++++++++++++ .../replays/test_ae_memory_revisions_get.py | 32 ++++++++++++++++ .../test_ae_memory_revisions_private_list.py | 32 ++++++++++++++++ 13 files changed, 434 insertions(+) create mode 100644 tests/unit/vertexai/genai/replays/test_ae_memories_delete.py create mode 100644 tests/unit/vertexai/genai/replays/test_ae_memories_get.py create mode 100644 tests/unit/vertexai/genai/replays/test_ae_memories_private_create.py create mode 100644 tests/unit/vertexai/genai/replays/test_ae_memories_private_generate.py create mode 100644 tests/unit/vertexai/genai/replays/test_ae_memories_private_get_generate_memories_operation.py create mode 100644 tests/unit/vertexai/genai/replays/test_ae_memories_private_get_memory_operation.py create mode 100644 tests/unit/vertexai/genai/replays/test_ae_memories_private_list.py create mode 100644 tests/unit/vertexai/genai/replays/test_ae_memories_private_purge.py create mode 100644 tests/unit/vertexai/genai/replays/test_ae_memories_private_retrieve.py create mode 100644 tests/unit/vertexai/genai/replays/test_ae_memories_private_rollback.py create mode 100644 tests/unit/vertexai/genai/replays/test_ae_memories_private_update.py create mode 100644 tests/unit/vertexai/genai/replays/test_ae_memory_revisions_get.py create mode 100644 tests/unit/vertexai/genai/replays/test_ae_memory_revisions_private_list.py diff --git a/tests/unit/vertexai/genai/replays/test_ae_memories_delete.py b/tests/unit/vertexai/genai/replays/test_ae_memories_delete.py new file mode 100644 index 0000000000..cf7dc1189f --- /dev/null +++ b/tests/unit/vertexai/genai/replays/test_ae_memories_delete.py @@ -0,0 +1,32 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# pylint: disable=protected-access,bad-continuation,missing-function-docstring + +from tests.unit.vertexai.genai.replays import pytest_helper +from vertexai._genai import types + + +def test_delete_memory(client): + ae_memory_operation = client.agent_engines.memories.delete( + name="projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/memories/5605466683931099136", + ) + assert isinstance(ae_memory_operation, types.DeleteAgentEngineMemoryOperation) + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), + test_method="agent_engines.memories.delete", +) diff --git a/tests/unit/vertexai/genai/replays/test_ae_memories_get.py b/tests/unit/vertexai/genai/replays/test_ae_memories_get.py new file mode 100644 index 0000000000..4b5fde1f8c --- /dev/null +++ b/tests/unit/vertexai/genai/replays/test_ae_memories_get.py @@ -0,0 +1,32 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# pylint: disable=protected-access,bad-continuation,missing-function-docstring + +from tests.unit.vertexai.genai.replays import pytest_helper +from vertexai._genai import types + + +def test_get_memory(client): + memory_name = "projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/memories/3858070028511346688" + ae_memory = client.agent_engines.memories.get(name=memory_name) + assert isinstance(ae_memory, types.Memory) + assert ae_memory.name == memory_name + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), + test_method="agent_engines.memories.get", +) diff --git a/tests/unit/vertexai/genai/replays/test_ae_memories_private_create.py b/tests/unit/vertexai/genai/replays/test_ae_memories_private_create.py new file mode 100644 index 0000000000..1aba8565e2 --- /dev/null +++ b/tests/unit/vertexai/genai/replays/test_ae_memories_private_create.py @@ -0,0 +1,34 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# pylint: disable=protected-access,bad-continuation,missing-function-docstring + +from tests.unit.vertexai.genai.replays import pytest_helper +from vertexai._genai import types + + +def test_private_create_memory(client): + ae_memory_operation = client.agent_engines.memories._create( + name="projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584", + fact="memory_fact", + scope={"user_id": "123"}, + ) + assert isinstance(ae_memory_operation, types.AgentEngineMemoryOperation) + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), + test_method="agent_engines.memories._create", +) diff --git a/tests/unit/vertexai/genai/replays/test_ae_memories_private_generate.py b/tests/unit/vertexai/genai/replays/test_ae_memories_private_generate.py new file mode 100644 index 0000000000..c577b20277 --- /dev/null +++ b/tests/unit/vertexai/genai/replays/test_ae_memories_private_generate.py @@ -0,0 +1,35 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# pylint: disable=protected-access,bad-continuation,missing-function-docstring + +from tests.unit.vertexai.genai.replays import pytest_helper +from vertexai._genai import types + + +def test_private_generate_memory(client): + ae_memory_operation = client.agent_engines.memories._generate( + name="projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584", + vertex_session_source=types.GenerateMemoriesRequestVertexSessionSource( + session="projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/sessions/6922431337672474624" + ), + ) + assert isinstance(ae_memory_operation, types.AgentEngineGenerateMemoriesOperation) + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), + test_method="agent_engines.memories._generate", +) diff --git a/tests/unit/vertexai/genai/replays/test_ae_memories_private_get_generate_memories_operation.py b/tests/unit/vertexai/genai/replays/test_ae_memories_private_get_generate_memories_operation.py new file mode 100644 index 0000000000..b888f606ef --- /dev/null +++ b/tests/unit/vertexai/genai/replays/test_ae_memories_private_get_generate_memories_operation.py @@ -0,0 +1,32 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# pylint: disable=protected-access,bad-continuation,missing-function-docstring + +from tests.unit.vertexai.genai.replays import pytest_helper +from vertexai._genai import types + + +def test_private_get_generate_memories_operation(client): + memory_operation = client.agent_engines.memories._get_generate_memories_operation( + operation_name="projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/operations/5669315676343369728" + ) + assert isinstance(memory_operation, types.AgentEngineGenerateMemoriesOperation) + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), + test_method="agent_engines.memories._get_generate_memories_operation", +) diff --git a/tests/unit/vertexai/genai/replays/test_ae_memories_private_get_memory_operation.py b/tests/unit/vertexai/genai/replays/test_ae_memories_private_get_memory_operation.py new file mode 100644 index 0000000000..587b1f7806 --- /dev/null +++ b/tests/unit/vertexai/genai/replays/test_ae_memories_private_get_memory_operation.py @@ -0,0 +1,32 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# pylint: disable=protected-access,bad-continuation,missing-function-docstring + +from tests.unit.vertexai.genai.replays import pytest_helper +from vertexai._genai import types + + +def test_private_get_memory_operation(client): + memory_operation = client.agent_engines.memories._get_memory_operation( + operation_name="projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/memories/3858070028511346688/operations/1044963283964002304" + ) + assert isinstance(memory_operation, types.AgentEngineMemoryOperation) + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), + test_method="agent_engines.memories._get_memory_operation", +) diff --git a/tests/unit/vertexai/genai/replays/test_ae_memories_private_list.py b/tests/unit/vertexai/genai/replays/test_ae_memories_private_list.py new file mode 100644 index 0000000000..0b564591c5 --- /dev/null +++ b/tests/unit/vertexai/genai/replays/test_ae_memories_private_list.py @@ -0,0 +1,32 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# pylint: disable=protected-access,bad-continuation,missing-function-docstring + +from tests.unit.vertexai.genai.replays import pytest_helper +from vertexai._genai import types + + +def test_private_list_memory(client): + ae_name = "projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584" + memory_list = client.agent_engines.memories._list(name=ae_name) + assert isinstance(memory_list, types.ListReasoningEnginesMemoriesResponse) + assert isinstance(memory_list.memories[0], types.Memory) + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), + test_method="agent_engines.memories._list", +) diff --git a/tests/unit/vertexai/genai/replays/test_ae_memories_private_purge.py b/tests/unit/vertexai/genai/replays/test_ae_memories_private_purge.py new file mode 100644 index 0000000000..bb4a4a473c --- /dev/null +++ b/tests/unit/vertexai/genai/replays/test_ae_memories_private_purge.py @@ -0,0 +1,35 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# pylint: disable=protected-access,bad-continuation,missing-function-docstring + +from tests.unit.vertexai.genai.replays import pytest_helper +from vertexai._genai import types + + +def test_private_purge(client): + ae_memory_purge_operation = client.agent_engines.memories._purge( + name="projects/964831358985/locations/us-central1/reasoningEngines/6086402690647064576", + filter="scope.user_id=123", + ) + assert isinstance( + ae_memory_purge_operation, types.AgentEnginePurgeMemoriesOperation + ) + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), + test_method="agent_engines.memories._purge", +) diff --git a/tests/unit/vertexai/genai/replays/test_ae_memories_private_retrieve.py b/tests/unit/vertexai/genai/replays/test_ae_memories_private_retrieve.py new file mode 100644 index 0000000000..9c7ebda52b --- /dev/null +++ b/tests/unit/vertexai/genai/replays/test_ae_memories_private_retrieve.py @@ -0,0 +1,38 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# pylint: disable=protected-access,bad-continuation,missing-function-docstring + +from tests.unit.vertexai.genai.replays import pytest_helper +from vertexai._genai import types + + +def test_private_retrieve(client): + ae_name = "projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584" + retrieved_memories = client.agent_engines.memories._retrieve( + name=ae_name, + scope={"user_id": "123"}, + ) + assert isinstance(retrieved_memories, types.RetrieveMemoriesResponse) + assert isinstance( + retrieved_memories.retrieved_memories[0], + types.RetrieveMemoriesResponseRetrievedMemory, + ) + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), + test_method="agent_engines.memories._retrieve", +) diff --git a/tests/unit/vertexai/genai/replays/test_ae_memories_private_rollback.py b/tests/unit/vertexai/genai/replays/test_ae_memories_private_rollback.py new file mode 100644 index 0000000000..e9ff84fdb1 --- /dev/null +++ b/tests/unit/vertexai/genai/replays/test_ae_memories_private_rollback.py @@ -0,0 +1,33 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# pylint: disable=protected-access,bad-continuation,missing-function-docstring + +from tests.unit.vertexai.genai.replays import pytest_helper +from vertexai._genai import types + + +def test_private_rollback(client): + rollback_operation = client.agent_engines.memories._rollback( + name="projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/memories/3858070028511346688", + target_revision_id="3001207491565453312", + ) + assert isinstance(rollback_operation, types.AgentEngineRollbackMemoryOperation) + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), + test_method="agent_engines.memories._rollback", +) diff --git a/tests/unit/vertexai/genai/replays/test_ae_memories_private_update.py b/tests/unit/vertexai/genai/replays/test_ae_memories_private_update.py new file mode 100644 index 0000000000..c11387bc36 --- /dev/null +++ b/tests/unit/vertexai/genai/replays/test_ae_memories_private_update.py @@ -0,0 +1,35 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# pylint: disable=protected-access,bad-continuation,missing-function-docstring + +from tests.unit.vertexai.genai.replays import pytest_helper +from vertexai._genai import types + + +def test_private_update_memory(client): + memory_name = "projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/memories/3858070028511346688" + memory_update_operation = client.agent_engines.memories._update( + name=memory_name, + fact="memory_fact_updated", + scope={"user_id": "123"}, + ) + assert isinstance(memory_update_operation, types.AgentEngineMemoryOperation) + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), + test_method="agent_engines.memories._update", +) diff --git a/tests/unit/vertexai/genai/replays/test_ae_memory_revisions_get.py b/tests/unit/vertexai/genai/replays/test_ae_memory_revisions_get.py new file mode 100644 index 0000000000..73d7ea049d --- /dev/null +++ b/tests/unit/vertexai/genai/replays/test_ae_memory_revisions_get.py @@ -0,0 +1,32 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# pylint: disable=protected-access,bad-continuation,missing-function-docstring + +from tests.unit.vertexai.genai.replays import pytest_helper +from vertexai._genai import types + + +def test_get_memory_revisions(client): + memory_name = "projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/memories/3858070028511346688/revisions/516064922187071488" + memory_revision = client.agent_engines.memories.revisions.get(name=memory_name) + assert isinstance(memory_revision, types.MemoryRevision) + assert memory_revision.name == memory_name + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), + test_method="agent_engines.memories.revisions.get", +) diff --git a/tests/unit/vertexai/genai/replays/test_ae_memory_revisions_private_list.py b/tests/unit/vertexai/genai/replays/test_ae_memory_revisions_private_list.py new file mode 100644 index 0000000000..a160745c28 --- /dev/null +++ b/tests/unit/vertexai/genai/replays/test_ae_memory_revisions_private_list.py @@ -0,0 +1,32 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# pylint: disable=protected-access,bad-continuation,missing-function-docstring + +from tests.unit.vertexai.genai.replays import pytest_helper +from vertexai._genai import types + + +def test_private_list_memory_revisions(client): + ae_name = "projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/memories/3858070028511346688" + memory_list = client.agent_engines.memories.revisions._list(name=ae_name) + assert isinstance(memory_list, types.ListAgentEngineMemoryRevisionsResponse) + assert isinstance(memory_list.memory_revisions[0], types.MemoryRevision) + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), + test_method="agent_engines.memories.revisions._list", +)