fix custom tests

Author: jsklan

tests/test_async_client.py

@@ -76,13 +76,6 @@ async def test_moved_fn(self) -> None:
             await self.co.list_connectors("dummy", dummy="dummy") # type: ignore
 
 
-    @unittest.skipIf(os.getenv("CO_API_URL") is not None, "Doesn't work in staging.")
-    async def test_generate(self) -> None:
-        response = await self.co.generate(
-            prompt='Please explain to me how LLMs work',
-        )
-        print(response)
-
     @unittest.skipIf(os.getenv("CO_API_URL") is not None, "Doesn't work in staging.")
     async def test_embed(self) -> None:
         response = await self.co.embed(

tests/test_bedrock_client.py

@@ -24,18 +24,20 @@ def _setup_boto3_env():
 @unittest.skipIf(None == os.getenv("TEST_AWS"), "tests skipped because TEST_AWS is not set")
 class TestClient(unittest.TestCase):
     platform: str = "bedrock"
-    client: cohere.AwsClient = cohere.BedrockClient(
-        aws_access_key=aws_access_key,
-        aws_secret_key=aws_secret_key,
-        aws_session_token=aws_session_token,
-        aws_region=aws_region,
-    )
     models: typing.Dict[str, str] = {
         "chat_model": "cohere.command-r-plus-v1:0",
         "embed_model": "cohere.embed-multilingual-v3",
         "generate_model": "cohere.command-text-v14",
     }
 
+    def setUp(self) -> None:
+        self.client = cohere.BedrockClient(
+            aws_access_key=aws_access_key,
+            aws_secret_key=aws_secret_key,
+            aws_session_token=aws_session_token,
+            aws_region=aws_region,
+        )
+
     def test_rerank(self) -> None:
         if self.platform != "sagemaker":
             self.skipTest("Only sagemaker supports rerank")
@@ -130,12 +132,13 @@ class TestBedrockClientV2(unittest.TestCase):
     since the request would fail with a signature mismatch otherwise.
     """
 
-    client: cohere.ClientV2 = cohere.BedrockClientV2(
-        aws_access_key=aws_access_key,
-        aws_secret_key=aws_secret_key,
-        aws_session_token=aws_session_token,
-        aws_region=aws_region,
-    )
+    def setUp(self) -> None:
+        self.client = cohere.BedrockClientV2(
+            aws_access_key=aws_access_key,
+            aws_secret_key=aws_secret_key,
+            aws_session_token=aws_session_token,
+            aws_region=aws_region,
+        )
 
     def test_embed(self) -> None:
         response = self.client.embed(

tests/test_client.py

@@ -72,13 +72,6 @@ def test_moved_fn(self) -> None:
         with self.assertRaises(ValueError):
             co.list_connectors("dummy", dummy="dummy")  # type: ignore
 
-    @unittest.skipIf(os.getenv("CO_API_URL") is not None, "Doesn't work in staging.")
-    def test_generate(self) -> None:
-        response = co.generate(
-            prompt='Please explain to me how LLMs work',
-        )
-        print(response)
-
     def test_embed(self) -> None:
         response = co.embed(
             texts=['hello', 'goodbye'],

tests/test_client_init.py

@@ -5,12 +5,21 @@
 import cohere
 from cohere import ToolMessage, UserMessage, AssistantMessage
 
+try:
+    import boto3
+    HAS_BOTO3 = True
+except ImportError:
+    HAS_BOTO3 = False
+
 class TestClientInit(unittest.TestCase):
-    def test_inits(self) -> None:
+    @unittest.skipUnless(HAS_BOTO3, "boto3 not installed")
+    def test_aws_inits(self) -> None:
         cohere.BedrockClient()
         cohere.BedrockClientV2()
         cohere.SagemakerClient()
         cohere.SagemakerClientV2()
+
+    def test_inits(self) -> None:
         cohere.Client(api_key="n/a")
         cohere.ClientV2(api_key="n/a")