mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-17 23:51:30 +00:00
378 lines
13 KiB
Python
378 lines
13 KiB
Python
import json
|
|
from collections.abc import Sequence
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
import litellm
|
|
from litellm.integrations.pointfive.upload_client import PointFiveUploadClient
|
|
from litellm.litellm_core_utils.url_utils import validate_url
|
|
from litellm.types.integrations.pointfive import PointFiveUploadFailure
|
|
|
|
API_URL = "https://api.pointfive.co/api/v1/ingestion"
|
|
UPLOAD_URL = "https://uploads.example.invalid/some/object.ndjson.gz?signature=sig"
|
|
OBJECT_KEY = "some/object.ndjson.gz"
|
|
BODY = b"gzipped-bytes"
|
|
|
|
|
|
def _presigned(status_code: int = 200) -> httpx.Response:
|
|
return _response(
|
|
status_code, {"uploadUrl": UPLOAD_URL, "objectKey": OBJECT_KEY, "expiresAt": "2026-08-25T14:35:00Z"}
|
|
)
|
|
|
|
|
|
def _response(status_code: int, payload: object) -> httpx.Response:
|
|
return httpx.Response(status_code, text=json.dumps(payload))
|
|
|
|
|
|
def _refused(status_code: int, error: str) -> httpx.Response:
|
|
"""The body PointFive sends with every refusal."""
|
|
return _response(status_code, {"success": False, "error": error})
|
|
|
|
|
|
def _accepted() -> httpx.Response:
|
|
return httpx.Response(200, text="")
|
|
|
|
|
|
def _no_content() -> httpx.Response:
|
|
return httpx.Response(204, text="")
|
|
|
|
|
|
class FakeHTTPClient:
|
|
"""
|
|
Stands in for AsyncHTTPHandler, including its habit of raising on error statuses.
|
|
|
|
Scripted results are consumed in order, and the last one repeats, so a test that
|
|
cares about a single behaviour passes a single result.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
presign: Sequence[httpx.Response | Exception] | None = None,
|
|
put: Sequence[httpx.Response | Exception] | None = None,
|
|
) -> None:
|
|
self.presign = list(presign) if presign else [_presigned()] # mutable-ok: results are consumed by popping
|
|
self.put_results = list(put) if put else [_accepted()] # mutable-ok: results are consumed by popping
|
|
self.presign_calls: list[dict] = []
|
|
self.put_calls: list[dict] = []
|
|
|
|
async def post(self, url, json=None, headers=None, **_):
|
|
self.presign_calls.append({"url": url, "json": json, "headers": headers or {}})
|
|
return _next_result(self.presign, url)
|
|
|
|
async def put(self, url, data=None, headers=None, follow_redirects=None, **_):
|
|
self.put_calls.append(
|
|
{"url": url, "data": data, "headers": headers or {}, "follow_redirects": follow_redirects}
|
|
)
|
|
return _next_result(self.put_results, url)
|
|
|
|
|
|
def _next_result(results: list, url: str) -> httpx.Response:
|
|
result = results.pop(0) if len(results) > 1 else results[0]
|
|
if isinstance(result, Exception):
|
|
raise result
|
|
if result.status_code >= 300:
|
|
request = httpx.Request("POST", url)
|
|
raise httpx.HTTPStatusError(
|
|
"boom",
|
|
request=request,
|
|
response=httpx.Response(result.status_code, text=result.text, headers=result.headers),
|
|
)
|
|
return result
|
|
|
|
|
|
async def _no_backoff(_seconds: float) -> None:
|
|
return None
|
|
|
|
|
|
def _trusting_validator(url: str) -> tuple[str, str]:
|
|
"""Stands in for validate_url so the fixture hosts need no DNS; the SSRF tests use the real one."""
|
|
return url, httpx.URL(url).host
|
|
|
|
|
|
def _client(
|
|
http_client: FakeHTTPClient,
|
|
max_retries: int = 3,
|
|
api_url: str = API_URL,
|
|
validate_upload_url=_trusting_validator,
|
|
) -> PointFiveUploadClient:
|
|
return PointFiveUploadClient(
|
|
api_key="p5tu_testkey",
|
|
api_url=api_url,
|
|
http_client=http_client,
|
|
max_retries=max_retries,
|
|
sleep=_no_backoff,
|
|
validate_upload_url=validate_upload_url,
|
|
)
|
|
|
|
|
|
def _presigned_for(upload_url: str) -> httpx.Response:
|
|
return _response(200, {"uploadUrl": upload_url, "objectKey": OBJECT_KEY, "expiresAt": "2026-08-25T14:35:00Z"})
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_uploads_the_body_to_the_url_the_api_returned():
|
|
http_client = FakeHTTPClient()
|
|
|
|
outcome = await _client(http_client).upload(BODY)
|
|
|
|
assert outcome == OBJECT_KEY
|
|
assert http_client.put_calls[0]["url"] == UPLOAD_URL
|
|
assert http_client.put_calls[0]["data"] == BODY
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_presign_request_is_authenticated_and_sized():
|
|
http_client = FakeHTTPClient()
|
|
|
|
await _client(http_client).upload(BODY)
|
|
|
|
call = http_client.presign_calls[0]
|
|
assert call["url"] == "https://api.pointfive.co/api/v1/ingestion/upload-url"
|
|
assert call["headers"]["Authorization"] == "Bearer p5tu_testkey"
|
|
assert call["json"] == {"kind": "LITELLM", "byteCount": len(BODY)}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_trailing_slash_on_the_api_url_is_tolerated():
|
|
"""A pasted URL often ends in a slash; it must not produce a double slash in the path."""
|
|
http_client = FakeHTTPClient()
|
|
|
|
await _client(http_client, api_url=API_URL + "/").upload(BODY)
|
|
|
|
assert http_client.presign_calls[0]["url"] == "https://api.pointfive.co/api/v1/ingestion/upload-url"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_no_bearer_token_is_sent_to_the_presigned_url():
|
|
"""The URL carries its own authorization, so the api key must not travel with it."""
|
|
http_client = FakeHTTPClient()
|
|
|
|
await _client(http_client).upload(BODY)
|
|
|
|
assert "Authorization" not in http_client.put_calls[0]["headers"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_the_upload_pins_the_host_and_never_follows_a_redirect():
|
|
http_client = FakeHTTPClient()
|
|
|
|
await _client(http_client).upload(BODY)
|
|
|
|
call = http_client.put_calls[0]
|
|
assert call["headers"]["Host"] == "uploads.example.invalid"
|
|
assert call["follow_redirects"] is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_redirected_upload_is_refused_rather_than_followed():
|
|
"""A presigned URL never redirects legitimately; following one is how a bad endpoint reaches inside."""
|
|
http_client = FakeHTTPClient(put=[httpx.Response(301, headers={"location": "http://169.254.169.254/"})])
|
|
|
|
outcome = await _client(http_client).upload(BODY)
|
|
|
|
assert outcome == PointFiveUploadFailure(
|
|
"presigned upload redirected with 301, refusing to follow", retryable=False
|
|
)
|
|
assert len(http_client.put_calls) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"upload_url",
|
|
[
|
|
"http://169.254.169.254/latest/meta-data",
|
|
"https://10.0.0.7/internal/bucket/object",
|
|
"http://127.0.0.1:9000/bucket/object",
|
|
],
|
|
)
|
|
async def test_an_upload_url_inside_the_network_is_refused_before_any_bytes_leave(upload_url):
|
|
http_client = FakeHTTPClient(presign=[_presigned_for(upload_url)])
|
|
|
|
outcome = await _client(http_client, validate_upload_url=validate_url).upload(BODY)
|
|
|
|
assert isinstance(outcome, PointFiveUploadFailure)
|
|
assert not outcome.retryable
|
|
assert outcome.detail.startswith("presigned upload url refused: ")
|
|
assert http_client.put_calls == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_an_operator_can_switch_destination_validation_off(monkeypatch):
|
|
"""litellm.user_url_validation is the proxy-wide switch every SSRF guard honours."""
|
|
monkeypatch.setattr(litellm, "user_url_validation", False)
|
|
http_client = FakeHTTPClient(presign=[_presigned_for("http://10.0.0.7/bucket/object")])
|
|
|
|
outcome = await _client(http_client, validate_upload_url=validate_url).upload(BODY)
|
|
|
|
assert outcome == OBJECT_KEY
|
|
assert http_client.put_calls[0]["url"] == "http://10.0.0.7/bucket/object"
|
|
assert "Host" not in http_client.put_calls[0]["headers"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_the_object_is_declared_as_gzipped_ndjson():
|
|
http_client = FakeHTTPClient()
|
|
|
|
await _client(http_client).upload(BODY)
|
|
|
|
assert http_client.put_calls[0]["headers"]["Content-Encoding"] == "gzip"
|
|
assert http_client.put_calls[0]["headers"]["Content-Type"] == "application/x-ndjson"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_each_retry_presigns_again():
|
|
"""A retry must never reuse a URL that was consumed or has expired."""
|
|
http_client = FakeHTTPClient(put=[httpx.Response(503), _accepted()])
|
|
|
|
outcome = await _client(http_client).upload(BODY)
|
|
|
|
assert outcome == OBJECT_KEY
|
|
assert len(http_client.presign_calls) == 2
|
|
assert len(http_client.put_calls) == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retryable_upload_failure_gives_up_after_max_retries():
|
|
http_client = FakeHTTPClient(put=[httpx.Response(503)])
|
|
|
|
outcome = await _client(http_client, max_retries=2).upload(BODY)
|
|
|
|
assert outcome == PointFiveUploadFailure("presigned upload returned 503, gave up after 2 attempts", retryable=True)
|
|
assert len(http_client.put_calls) == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_rejected_upload_is_not_retried():
|
|
http_client = FakeHTTPClient(put=[httpx.Response(403)])
|
|
|
|
outcome = await _client(http_client).upload(BODY)
|
|
|
|
assert outcome == PointFiveUploadFailure("presigned upload returned 403", retryable=False)
|
|
assert len(http_client.put_calls) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_bad_api_key_is_not_retried():
|
|
http_client = FakeHTTPClient(presign=[httpx.Response(401)])
|
|
|
|
outcome = await _client(http_client).upload(BODY)
|
|
|
|
assert outcome == PointFiveUploadFailure("pointfive api returned 401", retryable=False)
|
|
assert http_client.put_calls == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_the_reason_for_a_refusal_is_surfaced():
|
|
"""A 403 means the key no longer maps to an integration; the operator needs to read why."""
|
|
http_client = FakeHTTPClient(presign=[_refused(403, "no integration accepts uploads from this api key")])
|
|
|
|
outcome = await _client(http_client).upload(BODY)
|
|
|
|
assert outcome == PointFiveUploadFailure(
|
|
"pointfive api returned 403, no integration accepts uploads from this api key", retryable=False
|
|
)
|
|
assert http_client.put_calls == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_api_server_error_is_retried():
|
|
http_client = FakeHTTPClient(presign=[httpx.Response(503), _presigned()])
|
|
|
|
outcome = await _client(http_client).upload(BODY)
|
|
|
|
assert outcome == OBJECT_KEY
|
|
assert len(http_client.presign_calls) == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_too_many_requests_is_retried():
|
|
http_client = FakeHTTPClient(presign=[httpx.Response(429), _presigned()])
|
|
|
|
outcome = await _client(http_client).upload(BODY)
|
|
|
|
assert outcome == OBJECT_KEY
|
|
assert len(http_client.presign_calls) == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unreachable_api_is_retried_then_reported_as_retryable():
|
|
http_client = FakeHTTPClient(presign=(ConnectionError("down"),))
|
|
|
|
outcome = await _client(http_client, max_retries=2).upload(BODY)
|
|
|
|
assert isinstance(outcome, PointFiveUploadFailure)
|
|
assert outcome.retryable
|
|
assert "unreachable" in outcome.detail
|
|
assert len(http_client.presign_calls) == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_malformed_api_body_is_not_retried():
|
|
http_client = FakeHTTPClient(presign=[_response(200, {"objectKey": "k"})])
|
|
|
|
outcome = await _client(http_client).upload(BODY)
|
|
|
|
assert outcome == PointFiveUploadFailure("pointfive api returned an unreadable body", retryable=False)
|
|
assert http_client.put_calls == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_body_that_is_not_json_is_reported_as_unreadable():
|
|
http_client = FakeHTTPClient(presign=(httpx.Response(200, text="<html>gateway</html>"),))
|
|
|
|
outcome = await _client(http_client).upload(BODY)
|
|
|
|
assert outcome == PointFiveUploadFailure("pointfive api returned an unreadable body", retryable=False)
|
|
assert http_client.put_calls == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ping_reports_a_live_shipper():
|
|
http_client = FakeHTTPClient(presign=(_no_content(),))
|
|
|
|
failure = await _client(http_client).ping()
|
|
|
|
assert failure is None
|
|
assert http_client.presign_calls[0]["url"] == "https://api.pointfive.co/api/v1/ingestion/ping"
|
|
assert http_client.presign_calls[0]["json"] == {"kind": "LITELLM"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ping_surfaces_a_revoked_key():
|
|
http_client = FakeHTTPClient(presign=(_refused(403, "no integration accepts uploads from this api key"),))
|
|
|
|
failure = await _client(http_client).ping()
|
|
|
|
assert failure is not None
|
|
assert not failure.retryable
|
|
assert "no integration accepts uploads from this api key" in failure.detail
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ping_surfaces_an_unreachable_api():
|
|
http_client = FakeHTTPClient(presign=(ConnectionError("down"),))
|
|
|
|
failure = await _client(http_client).ping()
|
|
|
|
assert failure is not None
|
|
assert failure.retryable
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_transport_fault_on_the_upload_itself_is_retryable():
|
|
http_client = FakeHTTPClient(put=(ConnectionError("reset"),))
|
|
|
|
outcome = await _client(http_client, max_retries=1).upload(BODY)
|
|
|
|
assert isinstance(outcome, PointFiveUploadFailure)
|
|
assert outcome.retryable
|
|
assert "presigned upload unreachable" in outcome.detail
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_client_that_may_not_try_at_all_says_so():
|
|
"""max_upload_retries is validated as >= 1, so this guards the loop against a future zero."""
|
|
outcome = await _client(FakeHTTPClient(), max_retries=0).upload(BODY)
|
|
|
|
assert outcome == PointFiveUploadFailure("max_upload_retries must be at least 1", retryable=False)
|