chore: add unit tests for ModuleAgent template

yeesian · copybara-github · commit 16807d8f5ce9 · 2025-05-01T10:25:37.000-07:00
PiperOrigin-RevId: 753638103
diff --git a/tests/unit/vertex_langchain/test_agent_engine_templates_module.py b/tests/unit/vertex_langchain/test_agent_engine_templates_module.py
@@ -0,0 +1,89 @@
+# 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.
+#
+from vertexai import agent_engines
+from test_constants import test_agent
+
+_TEST_MODULE_NAME = "test_constants"
+_TEST_AGENT_NAME = "test_agent"
+_TEST_REGISTER_OPERATIONS = {"": ["query"], "stream": ["stream_query"]}
+_TEST_QUERY_INPUT = "test query"
+_TEST_STREAM_QUERY_INPUT = 5
+
+
+class TestModuleAgent:
+    def test_initialization(self):
+        agent = agent_engines.ModuleAgent(
+            module_name=_TEST_MODULE_NAME,
+            agent_name=_TEST_AGENT_NAME,
+            register_operations=_TEST_REGISTER_OPERATIONS,
+        )
+        assert agent._tmpl_attrs.get("module_name") == _TEST_MODULE_NAME
+        assert agent._tmpl_attrs.get("agent_name") == _TEST_AGENT_NAME
+        assert agent._tmpl_attrs.get("register_operations") == _TEST_REGISTER_OPERATIONS
+
+    def test_set_up(self):
+        agent = agent_engines.ModuleAgent(
+            module_name=_TEST_MODULE_NAME,
+            agent_name=_TEST_AGENT_NAME,
+            register_operations=_TEST_REGISTER_OPERATIONS,
+        )
+        assert agent._tmpl_attrs.get("agent") is None
+        agent.set_up()
+        assert agent._tmpl_attrs.get("agent") is not None
+
+    def test_clone(self):
+        agent = agent_engines.ModuleAgent(
+            module_name=_TEST_MODULE_NAME,
+            agent_name=_TEST_AGENT_NAME,
+            register_operations=_TEST_REGISTER_OPERATIONS,
+        )
+        agent.set_up()
+        assert agent._tmpl_attrs.get("agent") is not None
+        agent_clone = agent.clone()
+        assert agent._tmpl_attrs.get("agent") is not None
+        assert agent_clone._tmpl_attrs.get("agent") is None
+        agent_clone.set_up()
+        assert agent_clone._tmpl_attrs.get("agent") is not None
+
+    def test_query(self):
+        agent = agent_engines.ModuleAgent(
+            module_name=_TEST_MODULE_NAME,
+            agent_name=_TEST_AGENT_NAME,
+            register_operations=_TEST_REGISTER_OPERATIONS,
+        )
+        agent.set_up()
+        got_result = agent.query(input=_TEST_QUERY_INPUT)
+        expected_result = agent._tmpl_attrs.get("agent").query(input=_TEST_QUERY_INPUT)
+        assert got_result == expected_result
+        expected_result = test_agent.query(input=_TEST_QUERY_INPUT)
+        assert got_result == expected_result
+
+    def test_stream_query(self):
+        agent = agent_engines.ModuleAgent(
+            module_name=_TEST_MODULE_NAME,
+            agent_name=_TEST_AGENT_NAME,
+            register_operations=_TEST_REGISTER_OPERATIONS,
+        )
+        agent.set_up()
+        for got_result, expected_result in zip(
+            agent.stream_query(n=_TEST_STREAM_QUERY_INPUT),
+            agent._tmpl_attrs.get("agent").stream_query(n=_TEST_STREAM_QUERY_INPUT),
+        ):
+            assert got_result == expected_result
+        for got_result, expected_result in zip(
+            agent.stream_query(n=_TEST_STREAM_QUERY_INPUT),
+            test_agent.stream_query(n=_TEST_STREAM_QUERY_INPUT),
+        ):
+            assert got_result == expected_result
diff --git a/tests/unit/vertex_langchain/test_constants.py b/tests/unit/vertex_langchain/test_constants.py
@@ -0,0 +1,24 @@
+# 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.
+#
+class _CustomAgent:
+    def query(self, input: str):
+        return input
+
+    def stream_query(self, n: int):
+        for i in range(n):
+            yield i
+
+
+test_agent = _CustomAgent()
