117 lines
3.7 KiB
Python
117 lines
3.7 KiB
Python
"""
|
|
Test that async HTTP clients are properly cleaned up to prevent resource leaks.
|
|
Issue: https://github.com/BerriAI/litellm/issues/12107
|
|
"""
|
|
import asyncio
|
|
import os
|
|
import warnings
|
|
|
|
import pytest
|
|
|
|
import litellm
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_acompletion_resource_cleanup():
|
|
"""Test that acompletion doesn't leave unclosed client sessions."""
|
|
# Suppress warnings to check for them later
|
|
with warnings.catch_warnings(record=True) as w:
|
|
warnings.simplefilter("always")
|
|
|
|
# Make an async completion call
|
|
response = await litellm.acompletion(
|
|
model="gemini/gemini-2.0-flash-lite-001",
|
|
messages=[{"role": "user", "content": "Hello"}],
|
|
mock_response="Hi there! How can I help you today?",
|
|
)
|
|
|
|
# Check that response was received
|
|
assert (
|
|
response.choices[0].message.content == "Hi there! How can I help you today?"
|
|
)
|
|
|
|
# Manually close async clients
|
|
await litellm.close_litellm_async_clients()
|
|
|
|
# Give a small delay for any warnings to appear
|
|
await asyncio.sleep(0.1)
|
|
|
|
# Check for resource warnings
|
|
resource_warnings = [
|
|
warning
|
|
for warning in w
|
|
if "Unclosed" in str(warning.message)
|
|
and (
|
|
"client session" in str(warning.message)
|
|
or "connector" in str(warning.message)
|
|
)
|
|
]
|
|
|
|
# Should be no unclosed resource warnings
|
|
assert (
|
|
len(resource_warnings) == 0
|
|
), f"Found unclosed resources: {[str(w.message) for w in resource_warnings]}"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_multiple_acompletion_calls_cleanup():
|
|
"""Test that multiple acompletion calls reuse clients and don't leak resources."""
|
|
with warnings.catch_warnings(record=True) as w:
|
|
warnings.simplefilter("always")
|
|
|
|
# Make multiple async completion calls
|
|
for i in range(3):
|
|
response = await litellm.acompletion(
|
|
model="gemini/gemini-2.0-flash-lite-001",
|
|
messages=[{"role": "user", "content": f"Hello {i}"}],
|
|
mock_response=f"Response {i}",
|
|
)
|
|
assert response.choices[0].message.content == f"Response {i}"
|
|
|
|
# Clean up
|
|
await litellm.close_litellm_async_clients()
|
|
|
|
# Give a small delay for any warnings to appear
|
|
await asyncio.sleep(0.1)
|
|
|
|
# Check for resource warnings
|
|
resource_warnings = [
|
|
warning
|
|
for warning in w
|
|
if "Unclosed" in str(warning.message)
|
|
and (
|
|
"client session" in str(warning.message)
|
|
or "connector" in str(warning.message)
|
|
)
|
|
]
|
|
|
|
assert (
|
|
len(resource_warnings) == 0
|
|
), f"Found unclosed resources: {[str(w.message) for w in resource_warnings]}"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cleanup_function_is_safe_to_call_multiple_times():
|
|
"""Test that the cleanup function can be called multiple times safely."""
|
|
# This should not raise any errors
|
|
await litellm.close_litellm_async_clients()
|
|
await litellm.close_litellm_async_clients()
|
|
await litellm.close_litellm_async_clients()
|
|
|
|
# Should still work after multiple cleanups
|
|
response = await litellm.acompletion(
|
|
model="gemini/gemini-2.0-flash-lite-001",
|
|
messages=[{"role": "user", "content": "Hello"}],
|
|
mock_response="Hi!",
|
|
)
|
|
assert response.choices[0].message.content == "Hi!"
|
|
|
|
# Clean up again
|
|
await litellm.close_litellm_async_clients()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# Run the test
|
|
asyncio.run(test_acompletion_resource_cleanup())
|
|
print("✅ All tests passed!")
|