Cleanup async tests

Author: MarkDaoust

tests/test_discuss_async.py

@@ -12,52 +12,57 @@
 # 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 unittest.mock
 
-import asynctest
-from asynctest import mock as async_mock
+import sys
+import unittest
+
+if sys.version_info < (3, 11):
+  import asynctest
+  from asynctest import mock as async_mock
 
 import google.ai.generativelanguage as glm
 
 from google.generativeai import discuss
-from google.generativeai import client
-import google.generativeai as genai
 from absl.testing import absltest
 from absl.testing import parameterized
 
-# TODO: replace returns with 'assert' statements
+bases = (parameterized.TestCase,)
 
+if sys.version_info < (3, 11):
+  bases = bases + (asynctest.TestCase,)
 
-class AsyncTests(parameterized.TestCase, asynctest.TestCase):
-    async def test_chat_async(self):
-        client = async_mock.MagicMock()
+unittest.skipIf(sys.version_info >= (3,11), "asynctest is not suported on python 3.11+")
+class AsyncTests(*bases):
+    if sys.version_info < (3, 11):
+        async def test_chat_async(self):
+            client = async_mock.MagicMock()
 
-        observed_request = None
+            observed_request = None
 
-        async def fake_generate_message(
-            request: glm.GenerateMessageRequest,
-        ) -> glm.GenerateMessageResponse:
-            nonlocal observed_request
-            observed_request = request
-            return glm.GenerateMessageResponse(
-                candidates=[
-                    glm.Message(
-                        author="1", content="Why did the chicken cross the road?"
-                    )
-                ]
-            )
+            async def fake_generate_message(
+                request: glm.GenerateMessageRequest,
+            ) -> glm.GenerateMessageResponse:
+                nonlocal observed_request
+                observed_request = request
+                return glm.GenerateMessageResponse(
+                    candidates=[
+                        glm.Message(
+                            author="1", content="Why did the chicken cross the road?"
+                        )
+                    ]
+                )
 
-        client.generate_message = fake_generate_message
+            client.generate_message = fake_generate_message
 
-        observed_response = await discuss.chat_async(
-            model="models/bard",
-            context="Example Prompt",
-            examples=[["Example from human", "Example response from AI"]],
-            messages=["Tell me a joke"],
-            temperature=0.75,
-            candidate_count=1,
-            client=client,
-        )
+            observed_response = await discuss.chat_async(
+                model="models/bard",
+                context="Example Prompt",
+                examples=[["Example from human", "Example response from AI"]],
+                messages=["Tell me a joke"],
+                temperature=0.75,
+                candidate_count=1,
+                client=client,
+            )
 
         self.assertEqual(
             observed_request,

tests/test_text.py

@@ -17,9 +17,6 @@
 import unittest
 import unittest.mock as mock
 
-import asynctest
-from asynctest import mock as async_mock
-
 import google.ai.generativelanguage as glm
 
 from google.generativeai import text as text_service