test(http): type the query-preservation tests

This commit is contained in:
ryan-crabbe-berri 2026-10-30 21:55:05 -07:00
parent ecc977177d
commit 42a68c0b0b
2 changed files with 17 additions and 11 deletions

View file

@ -736,10 +736,12 @@ def _echo_query_transport() -> httpx.MockTransport:
@pytest.mark.parametrize("url,params,expected_query", QUERY_MERGE_CASES)
@pytest.mark.parametrize("stream", [False, True])
def test_sync_get_sends_url_query_bytes_unchanged(url, params, expected_query, stream):
handler = HTTPHandler(client=httpx.Client(transport=_echo_query_transport()))
def test_sync_get_sends_url_query_bytes_unchanged(
url: str, params: dict[str, str | int] | None, expected_query: bytes, stream: bool
):
handler: Final = HTTPHandler(client=httpx.Client(transport=_echo_query_transport()))
try:
response = handler.get(url, params=params, stream=stream)
response: Final = handler.get(url, params=params, stream=stream)
assert response.read() == expected_query
finally:
handler.client.close()
@ -749,12 +751,14 @@ def test_sync_get_sends_url_query_bytes_unchanged(url, params, expected_query, s
@pytest.mark.parametrize(
"get_kwargs", [{}, {"stream": True}, {"max_response_bytes": 1000}], ids=["plain", "stream", "bounded"]
)
async def test_async_get_sends_url_query_bytes_unchanged(url, params, expected_query, get_kwargs):
handler = AsyncHTTPHandler()
async def test_async_get_sends_url_query_bytes_unchanged(
url: str, params: dict[str, str | int] | None, expected_query: bytes, get_kwargs: dict[str, bool | int]
):
handler: Final = AsyncHTTPHandler()
await handler.client.aclose()
handler.client = httpx.AsyncClient(transport=_echo_query_transport())
try:
response = await handler.get(url, params=params, **get_kwargs)
response: Final = await handler.get(url, params=params, **get_kwargs)
assert await response.aread() == expected_query
finally:
await handler.client.aclose()

View file

@ -68,7 +68,7 @@ from litellm.responses.streaming_iterator import (
BaseResponsesAPIStreamingIterator,
MockResponsesAPIStreamingIterator,
)
from litellm.types.llms.openai import HttpxBinaryResponseContent, ResponsesAPIResponse
from litellm.types.llms.openai import FileContentRequest, HttpxBinaryResponseContent, ResponsesAPIResponse
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import (
ImageObject,
@ -3206,7 +3206,9 @@ SIGNED_DOWNLOAD_QUERY: Final = "X-Amz-Signature=abc&ResponseContentLanguage=&p=a
class _SignedDownloadFilesConfig(AnthropicFilesConfig):
def transform_file_content_request(self, file_content_request, optional_params, litellm_params):
def transform_file_content_request(
self, file_content_request: FileContentRequest, optional_params: dict, litellm_params: dict
) -> tuple[str, dict]:
return f"https://downloads.example/{file_content_request['file_id']}?{SIGNED_DOWNLOAD_QUERY}", {
"token": "from-params",
"extra": "1",
@ -3214,13 +3216,13 @@ class _SignedDownloadFilesConfig(AnthropicFilesConfig):
async def test_streaming_file_content_keeps_signed_url_query_and_appends_new_params():
async_client = AsyncHTTPHandler()
async_client: Final = AsyncHTTPHandler()
await async_client.client.aclose()
async_client.client = httpx.AsyncClient(
transport=httpx.MockTransport(lambda request: httpx.Response(200, request=request, content=request.url.query))
)
try:
result = await BaseLLMHTTPHandler().async_retrieve_file_content_streaming(
result: Final = await BaseLLMHTTPHandler().async_retrieve_file_content_streaming(
file_content_request={"file_id": "file-abc"},
provider_config=_SignedDownloadFilesConfig(),
litellm_params={"api_key": "sk-test"},
@ -3229,7 +3231,7 @@ async def test_streaming_file_content_keeps_signed_url_query_and_appends_new_par
chunk_size=1024,
client=async_client,
)
body = b"".join([chunk async for chunk in result.stream_iterator])
body: Final = b"".join([chunk async for chunk in result.stream_iterator])
finally:
await async_client.client.aclose()