Add main async code match test (#299)

* Remove f-string from doc

* Add async_code_match

* Remove async code match for separate modules

* blacken

* Remove looping for working_dirs

* fpath.name is same as split

* Move decorator checks to separate func

* blacken

* fix unbound param

Author: mayureshagashe2105

google/generativeai/answer.py

@@ -249,9 +249,9 @@ def generate_answer(
     client: glm.GenerativeServiceClient | None = None,
     request_options: dict[str, Any] | None = None,
 ):
-    f"""
+    """
     Calls the GenerateAnswer API and returns a `types.Answer` containing the response.
-    
+
     You can pass a literal list of text chunks:
 
     >>> from google.generativeai import answer
@@ -320,6 +320,7 @@ async def generate_answer_async(
     safety_settings: safety_types.SafetySettingOptions | None = None,
     temperature: float | None = None,
     client: glm.GenerativeServiceClient | None = None,
+    request_options: dict[str, Any] | None = None,
 ):
     """
     Calls the API and returns a `types.Answer` containing the answer.
@@ -341,6 +342,9 @@ async def generate_answer_async(
     Returns:
         A `types.Answer` containing the model's text answer response.
     """
+    if request_options is None:
+        request_options = {}
+
     if client is None:
         client = get_default_generative_async_client()
 
@@ -354,6 +358,6 @@ async def generate_answer_async(
         answer_style=answer_style,
     )
 
-    response = await client.generate_answer(request)
+    response = await client.generate_answer(request, **request_options)
 
     return response

tests/test_async_code_match.py

@@ -0,0 +1,106 @@
+# -*- coding: utf-8 -*-
+# Copyright 2024 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.
+
+import pathlib
+import ast
+import typing
+import re
+
+from absl.testing import absltest
+from absl.testing import parameterized
+
+EXEMPT_DIRS = ["notebook"]
+EXEMPT_DECORATORS = ["overload", "property", "setter", "abstractmethod", "staticmethod"]
+EXEMPT_FILES = ["client.py", "version.py", "discuss.py", "files.py"]
+EXEMPT_FUNCTIONS = ["to_dict", "_to_proto", "to_proto", "from_proto", "from_dict", "_from_dict"]
+
+
+class CodeMatch(absltest.TestCase):
+
+    def _maybe_trim_docstring(self, node):
+        if (
+            node.body
+            and isinstance(node.body[0], ast.Expr)
+            and isinstance(node.body[0].value, ast.Constant)
+        ):
+            node.body = node.body[1:]
+
+        return ast.unparse(node)
+
+    def _inspect_decorator_exemption(self, node, fpath) -> bool:
+        for decorator in node.decorator_list:
+            if isinstance(decorator, ast.Attribute):
+                if decorator.attr in EXEMPT_DECORATORS:
+                    return True
+            elif isinstance(decorator, ast.Name):
+                if decorator.id in EXEMPT_DECORATORS:
+                    return True
+            elif isinstance(decorator, ast.Call):
+                if decorator.func.attr in EXEMPT_DECORATORS:
+                    return True
+            else:
+                raise TypeError(
+                    f"Unknown decorator type {decorator}, during checking {node.name} from {fpath.name}"
+                )
+
+        return False
+
+    def _execute_code_match(self, source, asource):
+        asource = (
+            asource.replace("anext", "next")
+            .replace("aiter", "iter")
+            .replace("_async", "")
+            .replace("async ", "")
+            .replace("await ", "")
+            .replace("Async", "")
+            .replace("ASYNC_", "")
+        )
+        asource = re.sub(" *?# type: ignore", "", asource)
+        self.assertEqual(source, asource)
+
+    def test_code_match_for_async_methods(self):
+        for fpath in (pathlib.Path(__file__).parent.parent / "google").rglob("*.py"):
+            if fpath.name in EXEMPT_FILES or any([d in fpath.parts for d in EXEMPT_DIRS]):
+                continue
+            # print(f"Checking {fpath.absolute()}")
+            code_match_funcs: dict[str, ast.AST] = {}
+            source = fpath.read_text()
+            source_nodes = ast.parse(source)
+
+            for node in ast.walk(source_nodes):
+                if isinstance(
+                    node, (ast.FunctionDef, ast.AsyncFunctionDef)
+                ) and not node.name.startswith("_"):
+                    name = node.name[:-6] if node.name.endswith("_async") else node.name
+                    if name in EXEMPT_FUNCTIONS or self._inspect_decorator_exemption(node, fpath):
+                        continue
+                    # print(f"Checking {node.name}")
+
+                    if func_name := code_match_funcs.pop(name, None):
+                        snode, anode = (
+                            (func_name, node)
+                            if isinstance(node, ast.AsyncFunctionDef)
+                            else (node, func_name)
+                        )
+                        func_source = self._maybe_trim_docstring(snode)
+                        func_asource = self._maybe_trim_docstring(anode)
+                        self._execute_code_match(func_source, func_asource)
+                        # print(f"Matched {node.name}")
+                    else:
+                        code_match_funcs[node.name] = node
+
+
+if __name__ == "__main__":
+    absltest.main()

tests/test_embedding.py

@@ -179,34 +179,6 @@ def test_embed_content_called_with_request_options(self):
 
         self.client.embed_content.assert_called_once_with(request, **request_options)
 
-    @parameterized.named_parameters(
-        dict(
-            testcase_name="embedding.embed_content",
-            obj=embedding.embed_content,
-            aobj=embedding.embed_content_async,
-        ),
-    )
-    def test_async_code_match(self, obj, aobj):
-        import inspect
-        import re
-
-        source = inspect.getsource(obj)
-        asource = inspect.getsource(aobj)
-        source = re.sub('""".*"""', "", source, flags=re.DOTALL)
-        asource = re.sub('""".*"""', "", asource, flags=re.DOTALL)
-        asource = (
-            asource.replace("anext", "next")
-            .replace("aiter", "iter")
-            .replace("_async", "")
-            .replace("async ", "")
-            .replace("await ", "")
-            .replace("Async", "")
-            .replace("ASYNC_", "")
-        )
-
-        asource = re.sub(" *?# type: ignore", "", asource)
-        self.assertEqual(source, asource)
-
 
 if __name__ == "__main__":
     absltest.main()

tests/test_permission.py

@@ -173,59 +173,6 @@ def test_update_permission_failure_restricted_update_path(self):
                 {"grantee_type": permission_services.to_grantee_type("user")}
             )
 
-    @parameterized.named_parameters(
-        [
-            "create_permission",
-            retriever_services.Corpus.create_permission,
-            retriever_services.Corpus.create_permission_async,
-        ],
-        [
-            "list_permissions",
-            retriever_services.Corpus.list_permissions,
-            retriever_services.Corpus.list_permissions_async,
-        ],
-        [
-            "Permission.delete",
-            permission_services.Permission.delete,
-            permission_services.Permission.delete_async,
-        ],
-        [
-            "Permission.update",
-            permission_services.Permission.update,
-            permission_services.Permission.update_async,
-        ],
-        [
-            "Permission.get_permission",
-            permission_services.Permission.get,
-            permission_services.Permission.get_async,
-        ],
-        [
-            "permission.get_permission",
-            permission.get_permission,
-            permission.get_permission_async,
-        ],
-    )
-    def test_async_code_match(self, obj, aobj):
-        import inspect
-        import re
-
-        source = inspect.getsource(obj)
-        asource = inspect.getsource(aobj)
-        source = re.sub('""".*"""', "", source, flags=re.DOTALL)
-        asource = re.sub('""".*"""', "", asource, flags=re.DOTALL)
-        asource = (
-            asource.replace("anext", "next")
-            .replace("aiter", "iter")
-            .replace("_async", "")
-            .replace("async ", "")
-            .replace("await ", "")
-            .replace("Async", "")
-            .replace("ASYNC_", "")
-        )
-
-        asource = re.sub(" *?# type: ignore", "", asource)
-        self.assertEqual(source, asource)
-
     def test_create_corpus_called_with_request_options(self):
         self.client.create_corpus = unittest.mock.MagicMock()
         request = unittest.mock.ANY

tests/test_retriever.py

@@ -647,129 +647,6 @@ def test_batch_delete_chunks(self):
         delete_request = demo_document.batch_delete_chunks(chunks=[x.name, y.name])
         self.assertIsInstance(self.observed_requests[-1], glm.BatchDeleteChunksRequest)
 
-    @parameterized.named_parameters(
-        [
-            "create_corpus",
-            retriever.create_corpus,
-            retriever.create_corpus_async,
-        ],
-        [
-            "get_corpus",
-            retriever.get_corpus,
-            retriever.get_corpus_async,
-        ],
-        [
-            "delete_corpus",
-            retriever.delete_corpus,
-            retriever.delete_corpus_async,
-        ],
-        [
-            "list_corpora",
-            retriever.list_corpora,
-            retriever.list_corpora_async,
-        ],
-        [
-            "Corpus.create_document",
-            retriever_service.Corpus.create_document,
-            retriever_service.Corpus.create_document_async,
-        ],
-        [
-            "Corpus.get_document",
-            retriever_service.Corpus.get_document,
-            retriever_service.Corpus.get_document_async,
-        ],
-        [
-            "Corpus.update",
-            retriever_service.Corpus.update,
-            retriever_service.Corpus.update_async,
-        ],
-        [
-            "Corpus.query",
-            retriever_service.Corpus.query,
-            retriever_service.Corpus.query_async,
-        ],
-        [
-            "Corpus.list_documents",
-            retriever_service.Corpus.list_documents,
-            retriever_service.Corpus.list_documents_async,
-        ],
-        [
-            "Corpus.delete_document",
-            retriever_service.Corpus.delete_document,
-            retriever_service.Corpus.delete_document_async,
-        ],
-        [
-            "Document.create_chunk",
-            retriever_service.Document.create_chunk,
-            retriever_service.Document.create_chunk_async,
-        ],
-        [
-            "Document.get_chunk",
-            retriever_service.Document.get_chunk,
-            retriever_service.Document.get_chunk_async,
-        ],
-        [
-            "Document.batch_create_chunks",
-            retriever_service.Document.batch_create_chunks,
-            retriever_service.Document.batch_create_chunks_async,
-        ],
-        [
-            "Document.list_chunks",
-            retriever_service.Document.list_chunks,
-            retriever_service.Document.list_chunks_async,
-        ],
-        [
-            "Document.query",
-            retriever_service.Document.query,
-            retriever_service.Document.query_async,
-        ],
-        [
-            "Document.update",
-            retriever_service.Document.update,
-            retriever_service.Document.update_async,
-        ],
-        [
-            "Document.batch_update_chunks",
-            retriever_service.Document.batch_update_chunks,
-            retriever_service.Document.batch_update_chunks_async,
-        ],
-        [
-            "Document.delete_chunk",
-            retriever_service.Document.delete_chunk,
-            retriever_service.Document.delete_chunk_async,
-        ],
-        [
-            "Document.batch_delete_chunks",
-            retriever_service.Document.batch_delete_chunks,
-            retriever_service.Document.batch_delete_chunks_async,
-        ],
-        [
-            "Chunk.update",
-            retriever_service.Chunk.update,
-            retriever_service.Chunk.update_async,
-        ],
-    )
-    def test_async_code_match(self, obj, aobj):
-        import inspect
-        import re
-
-        source = inspect.getsource(obj)
-        asource = inspect.getsource(aobj)
-        source = re.sub('""".*"""', "", source, flags=re.DOTALL)
-        asource = re.sub('""".*"""', "", asource, flags=re.DOTALL)
-        asource = (
-            asource.replace("anext", "next")
-            .replace("aiter", "iter")
-            .replace("_async", "")
-            .replace("async ", "")
-            .replace("await ", "")
-            .replace("Async", "")
-            .replace("ASYNC_", "")
-        )
-
-        asource = re.sub(" *?# type: ignore", "", asource)
-        self.assertEqual(source, asource)
-
     @parameterized.parameters(
         {"method": "create_corpus"},
         {"method": "get_corpus"},