Update connection pooling tests to follow repository standards

- Use unittest.skipIf decorator for tests requiring API key
- Add tests that work without API key for basic verification
- Test httpx client creation with limits
- Test Cohere client initialization with pooling
- Follow existing test patterns in the repository
- All tests pass ruff and mypy checks

Author: Fede Kamelhar

tests/test_connection_pooling.py

@@ -1,3 +1,4 @@
+import os
 import time
 import unittest
 
@@ -9,138 +10,155 @@
 class TestConnectionPooling(unittest.TestCase):
     """Test suite for HTTP connection pooling functionality."""
 
-    def setUp(self):
-        """Set up test fixtures."""
-        self.client = cohere.ClientV2()
-
-    def test_client_has_connection_pooling_configured(self):
-        """Test that clients are configured with proper connection pooling limits."""
-        # Create clients
-        sync_client = cohere.Client()
-        v2_client = cohere.ClientV2()
-
-        # Verify they were created successfully
-        self.assertIsNotNone(sync_client)
-        self.assertIsNotNone(v2_client)
+    @classmethod
+    def setUpClass(cls):
+        """Set up class-level fixtures."""
+        # Check if API key is available for integration tests
+        cls.api_key_available = bool(os.environ.get("CO_API_KEY"))
+
+    def test_httpx_client_creation_with_limits(self):
+        """Test that httpx clients can be created with our connection pooling limits."""
+        # Test creating httpx client with limits (our implementation)
+        client_with_limits = httpx.Client(
+            timeout=300,
+            limits=httpx.Limits(
+                max_keepalive_connections=20,
+                max_connections=100,
+                keepalive_expiry=30.0,
+            ),
+        )
+        
+        # Verify the client was created successfully
+        self.assertIsNotNone(client_with_limits)
+        self.assertIsInstance(client_with_limits, httpx.Client)
+        
+        # The limits are applied internally - we can't directly access them
+        # but we verify the client works correctly with our configuration
+        
+        client_with_limits.close()
+
+    def test_cohere_client_initialization(self):
+        """Test that Cohere clients can be initialized with connection pooling."""
+        # Test with dummy API key - just verifies initialization works
+        try:
+            sync_client = cohere.Client(api_key="dummy-key")
+            v2_client = cohere.ClientV2(api_key="dummy-key")
+            
+            # Verify clients were created
+            self.assertIsNotNone(sync_client)
+            self.assertIsNotNone(v2_client)
+            
+        except Exception as e:
+            # Should not fail due to httpx configuration
+            if "httpx" in str(e).lower() or "limits" in str(e).lower():
+                self.fail(f"Failed to create client with connection pooling: {e}")
 
-        # The actual connection pooling is configured internally in httpx
-        # We mainly verify that clients can be created without errors
+    def test_custom_httpx_client_with_pooling(self):
+        """Test that custom httpx clients with connection pooling work correctly."""
+        # Create custom httpx client with explicit pooling configuration
+        custom_client = httpx.Client(
+            timeout=30,
+            limits=httpx.Limits(
+                max_keepalive_connections=10,
+                max_connections=50,
+                keepalive_expiry=20.0,
+            ),
+        )
 
-    def test_multiple_requests_performance(self):
-        """Test that multiple requests show performance improvement with connection pooling."""
-        # Skip if no API key is available
+        # Create Cohere client with custom httpx client
+        try:
+            client = cohere.ClientV2(api_key="dummy-key", httpx_client=custom_client)
+            self.assertIsNotNone(client)
+        finally:
+            custom_client.close()
+
+    def test_connection_pooling_vs_no_pooling_setup(self):
+        """Test creating clients with and without connection pooling."""
+        # Create httpx client without pooling
+        no_pool_httpx = httpx.Client(
+            timeout=30,
+            limits=httpx.Limits(
+                max_keepalive_connections=0,
+                max_connections=1,
+                keepalive_expiry=0,
+            ),
+        )
+        
+        # Verify both configurations work
         try:
-            client = cohere.ClientV2()
-        except Exception:
-            self.skipTest("No API key available")
+            pooled_client = cohere.ClientV2(api_key="dummy-key")
+            no_pool_client = cohere.ClientV2(api_key="dummy-key", httpx_client=no_pool_httpx)
+            
+            self.assertIsNotNone(pooled_client)
+            self.assertIsNotNone(no_pool_client)
+            
+        finally:
+            no_pool_httpx.close()
 
+    @unittest.skipIf(not os.environ.get("CO_API_KEY"), "API key not available")
+    def test_multiple_requests_performance(self):
+        """Test that multiple requests benefit from connection pooling."""
+        client = cohere.ClientV2()
+        
         response_times = []
-
-        # Make multiple requests and measure response times
+        
+        # Make multiple requests
         for i in range(3):
             start_time = time.time()
             try:
                 response = client.chat(
                     model="command-r-plus-08-2024",
-                    messages=[{"role": "user", "content": f"Say {i+1}"}],
+                    messages=[{"role": "user", "content": f"Say the number {i+1}"}],
                 )
                 elapsed = time.time() - start_time
                 response_times.append(elapsed)
-
-                # Verify response is valid
+                
+                # Verify response
                 self.assertIsNotNone(response)
                 self.assertIsNotNone(response.message)
-                self.assertIsNotNone(response.message.content)
-
-                # Wait to avoid rate limits
+                
+                # Rate limit protection
                 if i < 2:
                     time.sleep(2)
-
+                    
             except Exception as e:
-                # Skip test if rate limited or other API errors
                 if "429" in str(e) or "rate" in str(e).lower():
-                    self.skipTest(f"Rate limited: {e}")
+                    self.skipTest("Rate limited")
                 raise
-
-        # Generally, subsequent requests should be faster due to connection reuse
-        # We just verify all requests completed successfully
+        
+        # Verify all requests completed
         self.assertEqual(len(response_times), 3)
-        self.assertTrue(all(t > 0 for t in response_times))
+        
+        # Generally, subsequent requests should be faster due to connection reuse
+        # First request establishes connection, subsequent ones reuse it
+        print(f"Response times: {response_times}")
 
-    def test_streaming_with_connection_pooling(self):
+    @unittest.skipIf(not os.environ.get("CO_API_KEY"), "API key not available")
+    def test_streaming_with_pooling(self):
         """Test that streaming works correctly with connection pooling."""
-        try:
-            client = cohere.ClientV2()
-        except Exception:
-            self.skipTest("No API key available")
-
+        client = cohere.ClientV2()
+        
         try:
             response = client.chat_stream(
                 model="command-r-plus-08-2024",
                 messages=[{"role": "user", "content": "Count to 3"}],
             )
-
+            
             chunks = []
             for event in response:
                 if event.type == "content-delta":
                     chunks.append(event.delta.message.content.text)
-
-            # Verify we received chunks
+            
+            # Verify streaming worked
             self.assertGreater(len(chunks), 0)
             full_response = "".join(chunks)
             self.assertGreater(len(full_response), 0)
-
+            
         except Exception as e:
             if "429" in str(e) or "rate" in str(e).lower():
-                self.skipTest(f"Rate limited: {e}")
+                self.skipTest("Rate limited")
             raise
 
-    def test_custom_httpx_client_with_pooling(self):
-        """Test that custom httpx clients with connection pooling work correctly."""
-        # Create custom httpx client with explicit pooling configuration
-        custom_client = httpx.Client(
-            timeout=30,
-            limits=httpx.Limits(
-                max_keepalive_connections=10,
-                max_connections=50,
-                keepalive_expiry=20.0,
-            ),
-        )
-
-        # Create Cohere client with custom httpx client
-        client = cohere.ClientV2(httpx_client=custom_client)
-        self.assertIsNotNone(client)
-
-        # Clean up
-        custom_client.close()
-
-    def test_connection_pooling_vs_no_pooling(self):
-        """Compare performance with and without connection pooling."""
-        try:
-            # Test with connection pooling (default)
-            pooled_client = cohere.ClientV2()
-
-            # Test without connection pooling
-            no_pool_httpx = httpx.Client(
-                timeout=30,
-                limits=httpx.Limits(
-                    max_keepalive_connections=0,
-                    max_connections=1,
-                    keepalive_expiry=0,
-                ),
-            )
-            no_pool_client = cohere.ClientV2(httpx_client=no_pool_httpx)
-
-            # Just verify both clients work
-            self.assertIsNotNone(pooled_client)
-            self.assertIsNotNone(no_pool_client)
-
-            # Clean up
-            no_pool_httpx.close()
-
-        except Exception:
-            self.skipTest("No API key available")
-
 
 if __name__ == "__main__":
     unittest.main()