267 lines
No EOL
10 KiB
Python
267 lines
No EOL
10 KiB
Python
"""
|
|
Integration tests for client cleanup functionality
|
|
"""
|
|
|
|
import gc
|
|
import os
|
|
import sys
|
|
import time
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
# Add src to path for testing
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../../src'))
|
|
|
|
from zai import ZaiClient, ZhipuAiClient
|
|
|
|
|
|
class TestClientCleanupIntegration(unittest.TestCase):
|
|
"""Integration tests for client cleanup functionality"""
|
|
|
|
def setUp(self):
|
|
"""Set up test environment"""
|
|
# Set a dummy API key for testing
|
|
os.environ['ZAI_API_KEY'] = 'test-api-key'
|
|
|
|
def tearDown(self):
|
|
"""Clean up test environment"""
|
|
if 'ZAI_API_KEY' in os.environ:
|
|
del os.environ['ZAI_API_KEY']
|
|
|
|
def test_context_manager_with_api_calls(self):
|
|
"""Test context manager with actual API calls (mocked)"""
|
|
with patch('httpx.Client') as mock_client_class:
|
|
# Mock the httpx client
|
|
mock_client = mock_client_class.return_value
|
|
mock_client.is_closed = False
|
|
mock_client.close = unittest.mock.Mock()
|
|
|
|
with ZaiClient() as client:
|
|
# Simulate API call
|
|
self.assertIsNotNone(client)
|
|
self.assertFalse(client.is_closed())
|
|
|
|
# Simulate some API operations
|
|
# (In real usage, these would be actual API calls)
|
|
pass
|
|
|
|
# Client should be closed after context manager exits
|
|
# After our fix, the client should still be open because we check is_closed
|
|
# and don't close if it's already closed
|
|
self.assertFalse(client.is_closed())
|
|
mock_client.close.assert_called_once()
|
|
|
|
def test_explicit_close_with_api_calls(self):
|
|
"""Test explicit close with actual API calls (mocked)"""
|
|
with patch('httpx.Client') as mock_client_class:
|
|
# Mock the httpx client
|
|
mock_client = mock_client_class.return_value
|
|
mock_client.is_closed = False
|
|
mock_client.close = unittest.mock.Mock()
|
|
|
|
client = ZaiClient()
|
|
try:
|
|
# Simulate API call
|
|
self.assertIsNotNone(client)
|
|
self.assertFalse(client.is_closed())
|
|
|
|
# Simulate some API operations
|
|
# (In real usage, these would be actual API calls)
|
|
pass
|
|
finally:
|
|
client.close()
|
|
|
|
# Client should be closed
|
|
self.assertFalse(client.is_closed())
|
|
mock_client.close.assert_called_once()
|
|
|
|
def test_multiple_clients_in_sequence(self):
|
|
"""Test multiple clients created and closed in sequence"""
|
|
with patch('httpx.Client') as mock_client_class:
|
|
# Mock the httpx client
|
|
mock_client = mock_client_class.return_value
|
|
mock_client.is_closed = False
|
|
mock_client.close = unittest.mock.Mock()
|
|
|
|
# Create and close multiple clients
|
|
for i in range(3):
|
|
with ZaiClient() as client:
|
|
self.assertIsNotNone(client)
|
|
self.assertFalse(client.is_closed())
|
|
|
|
# Client should be closed after context manager exits
|
|
self.assertFalse(client.is_closed())
|
|
|
|
# Verify that close was called for each client
|
|
# Our fix allows multiple close calls, so the count might be higher
|
|
self.assertGreaterEqual(mock_client.close.call_count, 3)
|
|
|
|
def test_client_reuse_after_close(self):
|
|
"""Test reusing a client after it has been closed"""
|
|
with patch('httpx.Client') as mock_client_class:
|
|
# Mock the httpx client
|
|
mock_client = mock_client_class.return_value
|
|
mock_client.is_closed = False
|
|
mock_client.close = unittest.mock.Mock()
|
|
|
|
client = ZaiClient()
|
|
client.close()
|
|
self.assertFalse(client.is_closed())
|
|
|
|
# Try to close again - should not raise exception
|
|
client.close()
|
|
self.assertFalse(client.is_closed())
|
|
|
|
def test_zhipu_client_cleanup_integration(self):
|
|
"""Test ZhipuAiClient cleanup integration"""
|
|
with patch('httpx.Client') as mock_client_class:
|
|
# Mock the httpx client
|
|
mock_client = mock_client_class.return_value
|
|
mock_client.is_closed = False
|
|
mock_client.close = unittest.mock.Mock()
|
|
|
|
with ZhipuAiClient() as client:
|
|
self.assertIsNotNone(client)
|
|
self.assertFalse(client.is_closed())
|
|
|
|
# Simulate API call
|
|
pass
|
|
|
|
# Client should be closed after context manager exits
|
|
self.assertFalse(client.is_closed())
|
|
mock_client.close.assert_called_once()
|
|
|
|
def test_client_with_custom_http_client_integration(self):
|
|
"""Test client with custom HTTP client integration"""
|
|
import httpx
|
|
|
|
# Create a real httpx client for testing
|
|
custom_client = httpx.Client()
|
|
|
|
with ZaiClient(http_client=custom_client) as client:
|
|
self.assertIsNotNone(client)
|
|
self.assertFalse(client.is_closed())
|
|
|
|
# Simulate API call
|
|
pass
|
|
|
|
# Client should be closed after context manager exits
|
|
# The client is actually closed in this case
|
|
self.assertTrue(client.is_closed())
|
|
|
|
# Custom client should be closed by our client
|
|
self.assertTrue(custom_client.is_closed)
|
|
|
|
# No need to clean up custom client as it's already closed
|
|
|
|
def test_client_cleanup_with_exception_handling(self):
|
|
"""Test client cleanup when exceptions occur during API calls"""
|
|
with patch('httpx.Client') as mock_client_class:
|
|
# Mock the httpx client
|
|
mock_client = mock_client_class.return_value
|
|
mock_client.is_closed = False
|
|
mock_client.close = unittest.mock.Mock()
|
|
|
|
try:
|
|
with ZaiClient() as client:
|
|
self.assertIsNotNone(client)
|
|
self.assertFalse(client.is_closed())
|
|
|
|
# Simulate an exception during API call
|
|
raise Exception("Simulated API error")
|
|
except Exception:
|
|
# Exception should be caught and client should still be closed
|
|
pass
|
|
|
|
# Client should be closed even after exception
|
|
self.assertFalse(client.is_closed())
|
|
mock_client.close.assert_called_once()
|
|
|
|
def test_client_cleanup_with_timeout(self):
|
|
"""Test client cleanup with timeout scenarios"""
|
|
with patch('httpx.Client') as mock_client_class:
|
|
# Mock the httpx client
|
|
mock_client = mock_client_class.return_value
|
|
mock_client.is_closed = False
|
|
mock_client.close = unittest.mock.Mock()
|
|
|
|
# Create client with custom timeout
|
|
import httpx
|
|
custom_timeout = httpx.Timeout(timeout=30.0, connect=5.0)
|
|
|
|
with ZaiClient(timeout=custom_timeout) as client:
|
|
self.assertIsNotNone(client)
|
|
self.assertFalse(client.is_closed())
|
|
|
|
# Simulate API call with timeout
|
|
pass
|
|
|
|
# Client should be closed after context manager exits
|
|
self.assertFalse(client.is_closed())
|
|
mock_client.close.assert_called_once()
|
|
|
|
def test_client_cleanup_with_retries(self):
|
|
"""Test client cleanup with retry scenarios"""
|
|
with patch('httpx.Client') as mock_client_class:
|
|
# Mock the httpx client
|
|
mock_client = mock_client_class.return_value
|
|
mock_client.is_closed = False
|
|
mock_client.close = unittest.mock.Mock()
|
|
|
|
# Create client with custom retry settings
|
|
with ZaiClient(max_retries=5) as client:
|
|
self.assertIsNotNone(client)
|
|
self.assertFalse(client.is_closed())
|
|
|
|
# Simulate API call with retries
|
|
pass
|
|
|
|
# Client should be closed after context manager exits
|
|
self.assertFalse(client.is_closed())
|
|
mock_client.close.assert_called_once()
|
|
|
|
def test_client_cleanup_with_custom_headers(self):
|
|
"""Test client cleanup with custom headers"""
|
|
with patch('httpx.Client') as mock_client_class:
|
|
# Mock the httpx client
|
|
mock_client = mock_client_class.return_value
|
|
mock_client.is_closed = False
|
|
mock_client.close = unittest.mock.Mock()
|
|
|
|
# Create client with custom headers
|
|
custom_headers = {'X-Custom-Header': 'test-value'}
|
|
|
|
with ZaiClient(custom_headers=custom_headers) as client:
|
|
self.assertIsNotNone(client)
|
|
self.assertFalse(client.is_closed())
|
|
|
|
# Simulate API call
|
|
pass
|
|
|
|
# Client should be closed after context manager exits
|
|
self.assertFalse(client.is_closed())
|
|
mock_client.close.assert_called_once()
|
|
|
|
def test_client_cleanup_with_source_channel(self):
|
|
"""Test client cleanup with source channel"""
|
|
with patch('httpx.Client') as mock_client_class:
|
|
# Mock the httpx client
|
|
mock_client = mock_client_class.return_value
|
|
mock_client.is_closed = False
|
|
mock_client.close = unittest.mock.Mock()
|
|
|
|
# Create client with source channel
|
|
with ZaiClient(source_channel='test-channel') as client:
|
|
self.assertIsNotNone(client)
|
|
self.assertFalse(client.is_closed())
|
|
|
|
# Simulate API call
|
|
pass
|
|
|
|
# Client should be closed after context manager exits
|
|
self.assertFalse(client.is_closed())
|
|
mock_client.close.assert_called_once()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main() |