mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
6a0d03914c
* test: drop the cwd-relative sys.path.insert calls from the test suite
TQ003 stands at 1,077 across 1,058 files, and 1,015 of them are the same shape:
sys.path.insert(0, os.path.abspath("../..")) and its deeper siblings. The
argument resolves against the working directory rather than the file, so from
the repo root, where every job runs pytest, it inserts the directory two levels
above the checkout. It has never pointed at litellm. The package is installed
into the environment anyway, which is what actually makes the import work, and
what the rule's message has said all along.
Removing them leaves 1,634 imports of sys and os with no remaining reference,
and those go too, except where another test module imports the name back out of
the file. The rest of TQ003 is 62 call sites that resolve against __file__ or a
variable, which are a different question and are left alone.
Collection is identical either way: 45,871 tests and the same 51 pre-existing
collection errors before and after, and ruff reports no new undefined name.
* test: drop the duplicate imports the sys.path sweep exposed to F811
* test(pre-call-utils): restore the os import the new bedrock tests need
390 lines
11 KiB
Python
390 lines
11 KiB
Python
"""
|
|
Regression tests for streaming connection pool leak fix.
|
|
"""
|
|
|
|
import asyncio
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import anyio
|
|
import httpx
|
|
import pytest
|
|
|
|
|
|
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
|
from litellm.llms.custom_httpx.aiohttp_transport import (
|
|
AiohttpResponseStream,
|
|
LiteLLMAiohttpTransport,
|
|
)
|
|
|
|
|
|
# ── aiohttp transport layer tests ──────────────────────────────
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aiohttp_transport_response_uses_stream_not_content():
|
|
"""handle_async_request must use stream= so aclose() propagates to AiohttpResponseStream."""
|
|
|
|
class FakeSession:
|
|
closed = False
|
|
|
|
def __init__(self):
|
|
try:
|
|
self._loop = asyncio.get_running_loop()
|
|
except RuntimeError:
|
|
self._loop = None
|
|
|
|
def request(self, **kwargs):
|
|
class Resp:
|
|
status = 200
|
|
headers = {}
|
|
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, *args):
|
|
pass
|
|
|
|
@property
|
|
def content(self):
|
|
class C:
|
|
async def iter_chunked(self, size):
|
|
yield b"data"
|
|
|
|
return C()
|
|
|
|
return Resp()
|
|
|
|
transport = LiteLLMAiohttpTransport(client=lambda: FakeSession()) # type: ignore
|
|
response = await transport.handle_async_request(
|
|
httpx.Request("GET", "http://example.com")
|
|
)
|
|
|
|
assert isinstance(response.stream, AiohttpResponseStream)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aiohttp_response_stream_aclose_releases_connection():
|
|
"""AiohttpResponseStream.aclose() must call __aexit__ on the aiohttp response."""
|
|
aexit_called = False
|
|
|
|
class MockResponse:
|
|
status = 200
|
|
headers = {}
|
|
|
|
@property
|
|
def content(self):
|
|
class C:
|
|
async def iter_chunked(self, size):
|
|
yield b"data"
|
|
|
|
return C()
|
|
|
|
async def __aexit__(self, *args):
|
|
nonlocal aexit_called
|
|
aexit_called = True
|
|
|
|
stream = AiohttpResponseStream(MockResponse()) # type: ignore
|
|
await stream.aclose()
|
|
assert aexit_called
|
|
|
|
|
|
# ── CustomStreamWrapper.aclose() tests ─────────────────────────
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aclose_falls_back_to_close():
|
|
"""OpenAI's AsyncStream has close() but not aclose(). Must fall back."""
|
|
close_called = False
|
|
|
|
class FakeAsyncStream:
|
|
async def close(self):
|
|
nonlocal close_called
|
|
close_called = True
|
|
|
|
wrapper = CustomStreamWrapper(
|
|
completion_stream=FakeAsyncStream(),
|
|
model=None,
|
|
logging_obj=MagicMock(),
|
|
custom_llm_provider=None,
|
|
)
|
|
|
|
await wrapper.aclose()
|
|
assert close_called
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aclose_prefers_aclose_over_close():
|
|
"""When both aclose() and close() exist, aclose() should be preferred."""
|
|
aclose_called = False
|
|
close_called = False
|
|
|
|
class FakeStream:
|
|
async def aclose(self):
|
|
nonlocal aclose_called
|
|
aclose_called = True
|
|
|
|
async def close(self):
|
|
nonlocal close_called
|
|
close_called = True
|
|
|
|
wrapper = CustomStreamWrapper(
|
|
completion_stream=FakeStream(),
|
|
model=None,
|
|
logging_obj=MagicMock(),
|
|
custom_llm_provider=None,
|
|
)
|
|
|
|
await wrapper.aclose()
|
|
assert aclose_called
|
|
assert not close_called
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aclose_completes_under_cancellation():
|
|
"""aclose() must shield cleanup from CancelledError so streams actually close."""
|
|
aclose_completed = False
|
|
|
|
class SlowCloseStream:
|
|
async def aclose(self):
|
|
await anyio.sleep(0)
|
|
nonlocal aclose_completed
|
|
aclose_completed = True
|
|
|
|
wrapper = CustomStreamWrapper(
|
|
completion_stream=SlowCloseStream(),
|
|
model=None,
|
|
logging_obj=MagicMock(),
|
|
custom_llm_provider=None,
|
|
)
|
|
|
|
with anyio.CancelScope() as scope:
|
|
scope.cancel()
|
|
await wrapper.aclose()
|
|
|
|
assert aclose_completed
|
|
|
|
|
|
# ── Router stream_with_fallbacks cleanup tests ──────────────────
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stream_with_fallbacks_closes_stream_on_generator_close():
|
|
"""Closing the FallbackStreamWrapper must aclose() the underlying model_response
|
|
via stream_with_fallbacks' finally block."""
|
|
from litellm.router import Router
|
|
|
|
stream_closed = False
|
|
|
|
class FakeStream(CustomStreamWrapper):
|
|
def __init__(self):
|
|
super().__init__(
|
|
completion_stream=None,
|
|
model="test-model",
|
|
logging_obj=MagicMock(),
|
|
custom_llm_provider="openai",
|
|
)
|
|
self._items = ["chunk1", "chunk2", "chunk3"]
|
|
self._index = 0
|
|
|
|
def __aiter__(self):
|
|
return self
|
|
|
|
async def __anext__(self):
|
|
if self._index >= len(self._items):
|
|
raise StopAsyncIteration
|
|
item = self._items[self._index]
|
|
self._index += 1
|
|
return item
|
|
|
|
async def aclose(self):
|
|
nonlocal stream_closed
|
|
stream_closed = True
|
|
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "test-model",
|
|
"litellm_params": {
|
|
"model": "openai/test",
|
|
"api_key": "fake",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
fake_stream = FakeStream()
|
|
|
|
# Call _acompletion_streaming_iterator directly so we go through
|
|
# stream_with_fallbacks and its finally block
|
|
result = await router._acompletion_streaming_iterator(
|
|
model_response=fake_stream,
|
|
messages=[{"role": "user", "content": "hi"}],
|
|
initial_kwargs={"model": "test-model"},
|
|
)
|
|
|
|
# Consume one chunk then close (simulates client disconnect)
|
|
async for _ in result:
|
|
break
|
|
await result.aclose()
|
|
|
|
assert (
|
|
stream_closed
|
|
), "model_response stream was not closed by stream_with_fallbacks finally block"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stream_with_fallbacks_closes_stream_on_normal_completion():
|
|
"""stream_with_fallbacks must aclose() model_response even on normal completion."""
|
|
from litellm.router import Router
|
|
|
|
stream_closed = False
|
|
|
|
class FakeStream(CustomStreamWrapper):
|
|
def __init__(self):
|
|
super().__init__(
|
|
completion_stream=None,
|
|
model="test-model",
|
|
logging_obj=MagicMock(),
|
|
custom_llm_provider="openai",
|
|
)
|
|
self._items = ["chunk1"]
|
|
self._index = 0
|
|
|
|
def __aiter__(self):
|
|
return self
|
|
|
|
async def __anext__(self):
|
|
if self._index >= len(self._items):
|
|
raise StopAsyncIteration
|
|
item = self._items[self._index]
|
|
self._index += 1
|
|
return item
|
|
|
|
async def aclose(self):
|
|
nonlocal stream_closed
|
|
stream_closed = True
|
|
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "test-model",
|
|
"litellm_params": {
|
|
"model": "openai/test",
|
|
"api_key": "fake",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
fake_stream = FakeStream()
|
|
|
|
result = await router._acompletion_streaming_iterator(
|
|
model_response=fake_stream,
|
|
messages=[{"role": "user", "content": "hi"}],
|
|
initial_kwargs={"model": "test-model"},
|
|
)
|
|
|
|
# Exhaust the stream fully
|
|
async for _ in result:
|
|
pass
|
|
await result.aclose()
|
|
|
|
assert stream_closed, "model_response stream was not closed after normal completion"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stream_with_fallbacks_closes_both_on_fallback_disconnect():
|
|
"""When a fallback is triggered and the client disconnects during fallback
|
|
iteration, both model_response and fallback_response must be closed."""
|
|
from litellm.exceptions import MidStreamFallbackError
|
|
from litellm.router import Router
|
|
|
|
model_closed = False
|
|
fallback_closed = False
|
|
|
|
class FakeModelStream(CustomStreamWrapper):
|
|
"""Stream that raises MidStreamFallbackError immediately to trigger fallback."""
|
|
|
|
def __init__(self):
|
|
super().__init__(
|
|
completion_stream=None,
|
|
model="test-model",
|
|
logging_obj=MagicMock(),
|
|
custom_llm_provider="openai",
|
|
)
|
|
self.chunks = []
|
|
|
|
def __aiter__(self):
|
|
return self
|
|
|
|
async def __anext__(self):
|
|
raise MidStreamFallbackError(
|
|
message="test mid-stream error",
|
|
model="test-model",
|
|
llm_provider="openai",
|
|
generated_content="",
|
|
)
|
|
|
|
async def aclose(self):
|
|
nonlocal model_closed
|
|
model_closed = True
|
|
|
|
class FakeFallbackStream:
|
|
"""Fallback stream that yields chunks."""
|
|
|
|
def __init__(self):
|
|
self._items = ["fb1", "fb2", "fb3"]
|
|
self._index = 0
|
|
|
|
def __aiter__(self):
|
|
return self
|
|
|
|
async def __anext__(self):
|
|
if self._index >= len(self._items):
|
|
raise StopAsyncIteration
|
|
item = self._items[self._index]
|
|
self._index += 1
|
|
return item
|
|
|
|
async def aclose(self):
|
|
nonlocal fallback_closed
|
|
fallback_closed = True
|
|
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "test-model",
|
|
"litellm_params": {
|
|
"model": "openai/test",
|
|
"api_key": "fake",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
fake_model_stream = FakeModelStream()
|
|
fake_fallback_stream = FakeFallbackStream()
|
|
|
|
# Mock async_function_with_fallbacks_common_utils to return the fallback stream
|
|
# instead of actually calling through the full fallback machinery
|
|
with patch.object(
|
|
router,
|
|
"async_function_with_fallbacks_common_utils",
|
|
return_value=fake_fallback_stream,
|
|
):
|
|
result = await router._acompletion_streaming_iterator(
|
|
model_response=fake_model_stream,
|
|
messages=[{"role": "user", "content": "hi"}],
|
|
initial_kwargs={
|
|
"model": "test-model",
|
|
"fallbacks": ["other-model"],
|
|
},
|
|
)
|
|
|
|
# Consume one fallback chunk then close (simulates client disconnect)
|
|
async for _ in result:
|
|
break
|
|
await result.aclose()
|
|
|
|
assert model_closed, "model_response stream was not closed"
|
|
assert fallback_closed, "fallback_response stream was not closed"
|