mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-01 07:03:02 +00:00
6e0f0f7adb
* fix(s3_v2): bound each upload attempt, add jitter, retry on 429 and 503 Each S3 PUT now carries a per-attempt httpx timeout read from S3_UPLOAD_ATTEMPT_TIMEOUT_SECONDS (default 30s, clamped to 1..3600) instead of inheriting the handler's unbounded default. Retries back off with jitter instead of a fixed 1s then 2s, 429 joins 403, 500 and 503 in the in-call retry set, and transport errors or litellm.Timeout on an attempt are retried in the same call instead of ending it. Each flush logs one warning per batch of dropped uploads that names how many were terminal and how many aged out. * fix(s3_v2): retry uploads in a bounded loop and align tests with the new retry contract The async and sync retry helpers called themselves, which the recursive-function check rejects. They are now a bounded loop over the attempt count with the same outcome. The 429 integration test now expects in-call jittered backoff, and the aged-out test matches the new drop warning text. The sync transport test drives the real HTTPHandler so a dropped per-attempt timeout or a missed litellm.Timeout fails it, and the drop-reason test uses unequal counts so swapped labels fail. --------- Co-authored-by: yuneng-jiang <yuneng@berri.ai>
5850 lines
206 KiB
Python
5850 lines
206 KiB
Python
import asyncio, boto3, importlib, logging, os, pytest_asyncio
|
||
import copy
|
||
import json
|
||
import re
|
||
import sys
|
||
import textwrap
|
||
import threading
|
||
import time
|
||
import uuid
|
||
from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine
|
||
from contextlib import asynccontextmanager
|
||
from contextvars import Context
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from typing import Final
|
||
from unittest.mock import AsyncMock, MagicMock, patch
|
||
|
||
import httpx
|
||
import pytest
|
||
import respx
|
||
|
||
import litellm
|
||
from litellm.integrations.custom_logger import CustomLogger
|
||
from litellm.integrations.s3_v2 import S3BatchUploadError, S3Logger
|
||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||
from litellm.types.integrations.s3_v2 import s3BatchLoggingElement
|
||
from litellm.types.utils import StandardLoggingPayload
|
||
from collections import defaultdict
|
||
from litellm._logging import verbose_logger
|
||
from litellm.constants import LOGGING_WORKER_MAX_TIME_PER_COROUTINE, S3_UPLOAD_ATTEMPT_TIMEOUT_SECONDS
|
||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
|
||
|
||
_real_sleep: Final = asyncio.sleep
|
||
_NOW: Final = 1_000_000.0
|
||
|
||
|
||
def _assert_jittered_backoff(sleep_mock: MagicMock, retries: int) -> None:
|
||
"""Each retry waits a jittered exponential backoff: attempt n sleeps somewhere in [2**n / 2, 2**n * 1.5)."""
|
||
waits: Final = [sleep_call.args[0] for sleep_call in sleep_mock.call_args_list]
|
||
assert len(waits) == retries, waits
|
||
for attempt, wait in enumerate(waits):
|
||
assert (1 << attempt) * 0.5 <= wait < (1 << attempt) * 1.5, (attempt, wait)
|
||
|
||
|
||
class TestS3V2UnitTests:
|
||
"""Test that S3 v2 integration only uses safe_dumps and not json.dumps"""
|
||
|
||
def test_s3_v2_source_code_analysis(self):
|
||
"""Test that S3 v2 source code only imports and uses safe_dumps"""
|
||
import inspect
|
||
|
||
from litellm.integrations import s3_v2
|
||
|
||
# Get the source code of the s3_v2 module
|
||
source_code = inspect.getsource(s3_v2)
|
||
|
||
# Verify that json.dumps is not used directly in the code
|
||
assert "json.dumps(" not in source_code, "S3 v2 should not use json.dumps directly"
|
||
|
||
@patch("asyncio.create_task")
|
||
@patch("litellm.integrations.s3_v2.CustomBatchLogger.periodic_flush")
|
||
def test_s3_v2_endpoint_url(self, mock_periodic_flush, mock_create_task):
|
||
"""testing s3 endpoint url"""
|
||
from unittest.mock import AsyncMock, MagicMock
|
||
|
||
from litellm.types.integrations.s3_v2 import s3BatchLoggingElement
|
||
|
||
# Mock periodic_flush and create_task to prevent async task creation during init
|
||
mock_periodic_flush.return_value = None
|
||
mock_create_task.return_value = None
|
||
|
||
# Mock response for all tests
|
||
mock_response = MagicMock()
|
||
mock_response.status_code = 200
|
||
mock_response.raise_for_status = MagicMock()
|
||
|
||
# Create a test batch logging element
|
||
test_element = s3BatchLoggingElement(
|
||
s3_object_key="2025-09-14/test-key.json",
|
||
payload={"test": "data"},
|
||
s3_object_download_filename="test-file.json",
|
||
)
|
||
|
||
# Test 1: Custom endpoint URL with bucket name
|
||
s3_logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_endpoint_url="https://s3.amazonaws.com",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
|
||
s3_logger.async_httpx_client = AsyncMock()
|
||
s3_logger.async_httpx_client.put.return_value = mock_response
|
||
|
||
asyncio.run(s3_logger.async_upload_data_to_s3(test_element))
|
||
|
||
call_args = s3_logger.async_httpx_client.put.call_args
|
||
assert call_args is not None
|
||
url = call_args[0][0]
|
||
expected_url = "https://s3.amazonaws.com/test-bucket/2025-09-14/test-key.json"
|
||
assert url == expected_url, f"Expected URL {expected_url}, got {url}"
|
||
|
||
# Test 2: MinIO-compatible endpoint
|
||
s3_logger_minio = S3Logger(
|
||
s3_bucket_name="litellm-logs",
|
||
s3_endpoint_url="https://minio.example.com:9000",
|
||
s3_aws_access_key_id="minio-key",
|
||
s3_aws_secret_access_key="minio-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
|
||
s3_logger_minio.async_httpx_client = AsyncMock()
|
||
s3_logger_minio.async_httpx_client.put.return_value = mock_response
|
||
|
||
asyncio.run(s3_logger_minio.async_upload_data_to_s3(test_element))
|
||
|
||
call_args_minio = s3_logger_minio.async_httpx_client.put.call_args
|
||
assert call_args_minio is not None
|
||
url_minio = call_args_minio[0][0]
|
||
expected_minio_url = "https://minio.example.com:9000/litellm-logs/2025-09-14/test-key.json"
|
||
assert url_minio == expected_minio_url, f"Expected MinIO URL {expected_minio_url}, got {url_minio}"
|
||
|
||
# Test 3: Custom endpoint without bucket name (should fall back to default)
|
||
s3_logger_no_bucket = S3Logger(
|
||
s3_endpoint_url="https://s3.amazonaws.com",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
|
||
s3_logger_no_bucket.async_httpx_client = AsyncMock()
|
||
s3_logger_no_bucket.async_httpx_client.put.return_value = mock_response
|
||
|
||
asyncio.run(s3_logger_no_bucket.async_upload_data_to_s3(test_element))
|
||
|
||
call_args_no_bucket = s3_logger_no_bucket.async_httpx_client.put.call_args
|
||
assert call_args_no_bucket is not None
|
||
url_no_bucket = call_args_no_bucket[0][0]
|
||
# Should use default S3 URL format when bucket is missing (bucket becomes None in URL)
|
||
assert "s3.us-east-1.amazonaws.com" in url_no_bucket
|
||
assert "https://" in url_no_bucket
|
||
# Should not include the custom endpoint since bucket is missing
|
||
assert "https://s3.amazonaws.com/" not in url_no_bucket
|
||
|
||
# Test 4: Sync upload method with custom endpoint
|
||
s3_logger_sync = S3Logger(
|
||
s3_bucket_name="sync-bucket",
|
||
s3_endpoint_url="https://custom.s3.endpoint.com",
|
||
s3_aws_access_key_id="sync-key",
|
||
s3_aws_secret_access_key="sync-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
|
||
mock_sync_client = MagicMock()
|
||
mock_sync_client.put.return_value = mock_response
|
||
|
||
with patch(
|
||
"litellm.integrations.s3_v2.get_httpx_client",
|
||
return_value=mock_sync_client,
|
||
):
|
||
s3_logger_sync.upload_data_to_s3(test_element)
|
||
|
||
call_args_sync = mock_sync_client.put.call_args
|
||
assert call_args_sync is not None
|
||
url_sync = call_args_sync[0][0]
|
||
expected_sync_url = "https://custom.s3.endpoint.com/sync-bucket/2025-09-14/test-key.json"
|
||
assert url_sync == expected_sync_url, f"Expected sync URL {expected_sync_url}, got {url_sync}"
|
||
|
||
# Test 5: Download method with custom endpoint
|
||
s3_logger_download = S3Logger(
|
||
s3_bucket_name="download-bucket",
|
||
s3_endpoint_url="https://download.s3.endpoint.com",
|
||
s3_aws_access_key_id="download-key",
|
||
s3_aws_secret_access_key="download-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
|
||
mock_download_response = MagicMock()
|
||
mock_download_response.status_code = 200
|
||
mock_download_response.json = MagicMock(return_value={"downloaded": "data"})
|
||
s3_logger_download.async_httpx_client = AsyncMock()
|
||
s3_logger_download.async_httpx_client.get.return_value = mock_download_response
|
||
|
||
result = asyncio.run(s3_logger_download._download_object_from_s3("2025-09-14/download-test-key.json"))
|
||
|
||
call_args_download = s3_logger_download.async_httpx_client.get.call_args
|
||
assert call_args_download is not None
|
||
url_download = call_args_download[0][0]
|
||
expected_download_url = "https://download.s3.endpoint.com/download-bucket/2025-09-14/download-test-key.json"
|
||
assert url_download == expected_download_url, (
|
||
f"Expected download URL {expected_download_url}, got {url_download}"
|
||
)
|
||
|
||
assert result == {"downloaded": "data"}
|
||
|
||
@patch("asyncio.create_task")
|
||
@patch("litellm.integrations.s3_v2.CustomBatchLogger.periodic_flush")
|
||
def test_s3_v2_virtual_hosted_style(self, mock_periodic_flush, mock_create_task):
|
||
"""Test s3_use_virtual_hosted_style parameter for virtual-hosted-style URLs"""
|
||
from unittest.mock import AsyncMock, MagicMock
|
||
|
||
from litellm.types.integrations.s3_v2 import s3BatchLoggingElement
|
||
|
||
# Mock periodic_flush and create_task to prevent async task creation during init
|
||
mock_periodic_flush.return_value = None
|
||
mock_create_task.return_value = None
|
||
|
||
# Mock response for all tests
|
||
mock_response = MagicMock()
|
||
mock_response.status_code = 200
|
||
mock_response.raise_for_status = MagicMock()
|
||
|
||
# Create a test batch logging element
|
||
test_element = s3BatchLoggingElement(
|
||
s3_object_key="2025-09-14/test-key.json",
|
||
payload={"test": "data"},
|
||
s3_object_download_filename="test-file.json",
|
||
)
|
||
|
||
# Test 1: Virtual-hosted-style with custom endpoint
|
||
s3_logger_virtual = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_endpoint_url="https://s3.custom-endpoint.com",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_use_virtual_hosted_style=True,
|
||
)
|
||
|
||
s3_logger_virtual.async_httpx_client = AsyncMock()
|
||
s3_logger_virtual.async_httpx_client.put.return_value = mock_response
|
||
|
||
asyncio.run(s3_logger_virtual.async_upload_data_to_s3(test_element))
|
||
|
||
call_args = s3_logger_virtual.async_httpx_client.put.call_args
|
||
assert call_args is not None
|
||
url = call_args[0][0]
|
||
expected_url = "https://test-bucket.s3.custom-endpoint.com/2025-09-14/test-key.json"
|
||
assert url == expected_url, f"Expected virtual-hosted-style URL {expected_url}, got {url}"
|
||
|
||
# Test 2: Path-style (default behavior with s3_use_virtual_hosted_style=False)
|
||
s3_logger_path = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_endpoint_url="https://s3.custom-endpoint.com",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_use_virtual_hosted_style=False,
|
||
)
|
||
|
||
s3_logger_path.async_httpx_client = AsyncMock()
|
||
s3_logger_path.async_httpx_client.put.return_value = mock_response
|
||
|
||
asyncio.run(s3_logger_path.async_upload_data_to_s3(test_element))
|
||
|
||
call_args_path = s3_logger_path.async_httpx_client.put.call_args
|
||
assert call_args_path is not None
|
||
url_path = call_args_path[0][0]
|
||
expected_path_url = "https://s3.custom-endpoint.com/test-bucket/2025-09-14/test-key.json"
|
||
assert url_path == expected_path_url, f"Expected path-style URL {expected_path_url}, got {url_path}"
|
||
|
||
# Test 3: Virtual-hosted-style with http protocol
|
||
s3_logger_http = S3Logger(
|
||
s3_bucket_name="http-bucket",
|
||
s3_endpoint_url="http://minio.local:9000",
|
||
s3_aws_access_key_id="minio-key",
|
||
s3_aws_secret_access_key="minio-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_use_virtual_hosted_style=True,
|
||
)
|
||
|
||
s3_logger_http.async_httpx_client = AsyncMock()
|
||
s3_logger_http.async_httpx_client.put.return_value = mock_response
|
||
|
||
asyncio.run(s3_logger_http.async_upload_data_to_s3(test_element))
|
||
|
||
call_args_http = s3_logger_http.async_httpx_client.put.call_args
|
||
assert call_args_http is not None
|
||
url_http = call_args_http[0][0]
|
||
expected_http_url = "http://http-bucket.minio.local:9000/2025-09-14/test-key.json"
|
||
assert url_http == expected_http_url, (
|
||
f"Expected virtual-hosted-style URL with http {expected_http_url}, got {url_http}"
|
||
)
|
||
|
||
# Test 4: Sync upload method with virtual-hosted-style
|
||
s3_logger_sync_virtual = S3Logger(
|
||
s3_bucket_name="sync-bucket",
|
||
s3_endpoint_url="https://storage.example.com",
|
||
s3_aws_access_key_id="sync-key",
|
||
s3_aws_secret_access_key="sync-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_use_virtual_hosted_style=True,
|
||
)
|
||
|
||
mock_sync_client = MagicMock()
|
||
mock_sync_client.put.return_value = mock_response
|
||
|
||
with patch(
|
||
"litellm.integrations.s3_v2.get_httpx_client",
|
||
return_value=mock_sync_client,
|
||
):
|
||
s3_logger_sync_virtual.upload_data_to_s3(test_element)
|
||
|
||
call_args_sync = mock_sync_client.put.call_args
|
||
assert call_args_sync is not None
|
||
url_sync = call_args_sync[0][0]
|
||
expected_sync_url = "https://sync-bucket.storage.example.com/2025-09-14/test-key.json"
|
||
assert url_sync == expected_sync_url, (
|
||
f"Expected virtual-hosted-style sync URL {expected_sync_url}, got {url_sync}"
|
||
)
|
||
|
||
# Test 5: Download method with virtual-hosted-style
|
||
s3_logger_download_virtual = S3Logger(
|
||
s3_bucket_name="download-bucket",
|
||
s3_endpoint_url="https://download.endpoint.com",
|
||
s3_aws_access_key_id="download-key",
|
||
s3_aws_secret_access_key="download-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_use_virtual_hosted_style=True,
|
||
)
|
||
|
||
mock_download_response = MagicMock()
|
||
mock_download_response.status_code = 200
|
||
mock_download_response.json = MagicMock(return_value={"downloaded": "data"})
|
||
s3_logger_download_virtual.async_httpx_client = AsyncMock()
|
||
s3_logger_download_virtual.async_httpx_client.get.return_value = mock_download_response
|
||
|
||
result = asyncio.run(s3_logger_download_virtual._download_object_from_s3("2025-09-14/download-test-key.json"))
|
||
|
||
call_args_download = s3_logger_download_virtual.async_httpx_client.get.call_args
|
||
assert call_args_download is not None
|
||
url_download = call_args_download[0][0]
|
||
expected_download_url = "https://download-bucket.download.endpoint.com/2025-09-14/download-test-key.json"
|
||
assert url_download == expected_download_url, (
|
||
f"Expected virtual-hosted-style download URL {expected_download_url}, got {url_download}"
|
||
)
|
||
|
||
assert result == {"downloaded": "data"}
|
||
|
||
@patch("asyncio.create_task")
|
||
@patch("litellm.integrations.s3_v2.CustomBatchLogger.periodic_flush")
|
||
def test_s3_v2_put_url_encodes_spaces_in_object_key(self, mock_periodic_flush, mock_create_task):
|
||
from unittest.mock import AsyncMock
|
||
|
||
import requests
|
||
|
||
from litellm.types.integrations.s3_v2 import s3BatchLoggingElement
|
||
|
||
mock_periodic_flush.return_value = None
|
||
mock_create_task.return_value = None
|
||
|
||
mock_response = MagicMock()
|
||
mock_response.status_code = 200
|
||
mock_response.raise_for_status = MagicMock()
|
||
|
||
s3_object_key = "My Team/2025-09-14/test-key.json"
|
||
test_element = s3BatchLoggingElement(
|
||
s3_object_key=s3_object_key,
|
||
payload={"test": "data"},
|
||
s3_object_download_filename="test-file.json",
|
||
)
|
||
|
||
s3_logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_endpoint_url="https://s3.amazonaws.com",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
s3_logger.async_httpx_client = AsyncMock()
|
||
s3_logger.async_httpx_client.put.return_value = mock_response
|
||
|
||
asyncio.run(s3_logger.async_upload_data_to_s3(test_element))
|
||
|
||
call_args = s3_logger.async_httpx_client.put.call_args
|
||
assert call_args is not None
|
||
actual_url = call_args[0][0]
|
||
raw_url = f"https://s3.amazonaws.com/test-bucket/{s3_object_key}"
|
||
expected_url = requests.Request("PUT", raw_url).prepare().url
|
||
assert actual_url == expected_url
|
||
assert " " not in actual_url
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_upload_retries_on_s3_503():
|
||
"""
|
||
Test that async_upload_data_to_s3 retries on transient S3 503 Slow Down
|
||
and succeeds on the second attempt.
|
||
"""
|
||
from unittest.mock import AsyncMock, MagicMock
|
||
|
||
from litellm.types.integrations.s3_v2 import s3BatchLoggingElement
|
||
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
|
||
test_element = s3BatchLoggingElement(
|
||
s3_object_key="2025-09-14/test-retry.json",
|
||
payload={"test": "retry"},
|
||
s3_object_download_filename="test-retry.json",
|
||
)
|
||
|
||
# First call returns 503, second call returns 200
|
||
response_503 = MagicMock()
|
||
response_503.status_code = 503
|
||
response_503.text = ""
|
||
response_200 = MagicMock()
|
||
response_200.status_code = 200
|
||
response_200.text = ""
|
||
response_200.raise_for_status = MagicMock()
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = AsyncMock(side_effect=[response_503, response_200])
|
||
|
||
with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep:
|
||
await logger.async_upload_data_to_s3(test_element)
|
||
|
||
# Verify PUT was called twice (retry after 503)
|
||
assert logger.async_httpx_client.put.call_count == 2
|
||
_assert_jittered_backoff(mock_sleep, 1)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_upload_retries_on_s3_500():
|
||
"""
|
||
Test that async_upload_data_to_s3 retries on transient S3 500 errors.
|
||
"""
|
||
from unittest.mock import AsyncMock, MagicMock
|
||
|
||
from litellm.types.integrations.s3_v2 import s3BatchLoggingElement
|
||
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
|
||
test_element = s3BatchLoggingElement(
|
||
s3_object_key="2025-09-14/test-retry-500.json",
|
||
payload={"test": "retry-500"},
|
||
s3_object_download_filename="test-retry-500.json",
|
||
)
|
||
|
||
response_500 = MagicMock()
|
||
response_500.status_code = 500
|
||
response_500.text = ""
|
||
response_200 = MagicMock()
|
||
response_200.status_code = 200
|
||
response_200.text = ""
|
||
response_200.raise_for_status = MagicMock()
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = AsyncMock(side_effect=[response_500, response_200])
|
||
|
||
with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep:
|
||
await logger.async_upload_data_to_s3(test_element)
|
||
|
||
assert logger.async_httpx_client.put.call_count == 2
|
||
_assert_jittered_backoff(mock_sleep, 1)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_upload_exhausts_retries_on_persistent_503():
|
||
"""
|
||
Test that async_upload_data_to_s3 raises after exhausting all retries
|
||
on persistent S3 503.
|
||
"""
|
||
from unittest.mock import AsyncMock, MagicMock
|
||
|
||
from litellm.types.integrations.s3_v2 import s3BatchLoggingElement
|
||
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
|
||
test_element = s3BatchLoggingElement(
|
||
s3_object_key="2025-09-14/test-exhaust.json",
|
||
payload={"test": "exhaust"},
|
||
s3_object_download_filename="test-exhaust.json",
|
||
)
|
||
|
||
# All 3 attempts return 503
|
||
response_503 = MagicMock()
|
||
response_503.status_code = 503
|
||
response_503.text = ""
|
||
response_503.raise_for_status = MagicMock(side_effect=Exception("503 Service Unavailable"))
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = AsyncMock(return_value=response_503)
|
||
|
||
with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep:
|
||
with patch.object(logger, "handle_callback_failure") as mock_failure:
|
||
await logger.async_upload_data_to_s3(test_element)
|
||
|
||
# 3 PUT attempts total
|
||
assert logger.async_httpx_client.put.call_count == 3
|
||
# 2 sleeps (between attempts 1-2 and 2-3)
|
||
assert mock_sleep.call_count == 2
|
||
# Callback failure handler called after exhausting retries
|
||
mock_failure.assert_called_once_with(callback_name="S3Logger")
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_upload_retries_400_with_an_unknown_error_code():
|
||
"""
|
||
A 400 is outside the retry set, so an unknown <Code> gets a single PUT and the "retry" outcome
|
||
for the flush-level requeue, never an in-call backoff.
|
||
"""
|
||
from unittest.mock import AsyncMock, MagicMock
|
||
|
||
from litellm.types.integrations.s3_v2 import s3BatchLoggingElement
|
||
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
|
||
test_element = s3BatchLoggingElement(
|
||
s3_object_key="2025-09-14/test-retry-400.json",
|
||
payload={"test": "retry-400"},
|
||
s3_object_download_filename="test-retry-400.json",
|
||
)
|
||
|
||
response_400 = MagicMock()
|
||
response_400.status_code = 400
|
||
response_400.text = "<Error><Code>SomethingElse</Code></Error>"
|
||
response_400.raise_for_status = MagicMock(side_effect=Exception("400 Bad Request"))
|
||
response_200 = MagicMock()
|
||
response_200.status_code = 200
|
||
response_200.text = ""
|
||
response_200.raise_for_status = MagicMock()
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = AsyncMock(side_effect=[response_400, response_200])
|
||
|
||
with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep:
|
||
outcome = await logger.async_upload_data_to_s3(test_element)
|
||
|
||
assert logger.async_httpx_client.put.call_count == 1
|
||
mock_sleep.assert_not_awaited()
|
||
assert outcome is False
|
||
|
||
|
||
_SIGV4_ACCESS_KEY = re.compile(r"Credential=(AKIA\d+)/")
|
||
|
||
|
||
@pytest.fixture
|
||
def rotating_profile(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> str:
|
||
"""
|
||
A real botocore profile whose credential_process hands out a new key generation on every call and
|
||
expires inside the advisory refresh window, so RefreshableCredentials re-runs it on every property read.
|
||
"""
|
||
counter = tmp_path / "generation"
|
||
script = tmp_path / "rotate_credentials.py"
|
||
script.write_text(
|
||
textwrap.dedent(
|
||
f"""
|
||
import json, sys
|
||
from datetime import datetime, timedelta, timezone
|
||
from pathlib import Path
|
||
|
||
counter = Path({str(counter)!r})
|
||
generation = int(counter.read_text()) if counter.exists() else 0
|
||
counter.write_text(str(generation + 1))
|
||
expiry = (datetime.now(timezone.utc) + timedelta(minutes=12)).strftime("%Y-%m-%dT%H:%M:%SZ")
|
||
json.dump(
|
||
{{
|
||
"Version": 1,
|
||
"AccessKeyId": f"AKIA{{generation}}",
|
||
"SecretAccessKey": f"secret-{{generation}}",
|
||
"SessionToken": f"token-{{generation}}",
|
||
"Expiration": expiry,
|
||
}},
|
||
sys.stdout,
|
||
)
|
||
"""
|
||
)
|
||
)
|
||
profile = f"rotating-{uuid.uuid4().hex}"
|
||
(tmp_path / "config").write_text(f"[profile {profile}]\ncredential_process = {sys.executable} {script}\n")
|
||
monkeypatch.setenv("AWS_CONFIG_FILE", str(tmp_path / "config"))
|
||
return profile
|
||
|
||
|
||
def _generation(request: httpx.Request) -> tuple[str, str]:
|
||
"""(access key generation, session token generation) SigV4 baked into one request."""
|
||
access_key = _SIGV4_ACCESS_KEY.search(request.headers["Authorization"])
|
||
assert access_key is not None
|
||
return access_key.group(1).removeprefix("AKIA"), request.headers["X-Amz-Security-Token"].removeprefix("token-")
|
||
|
||
|
||
@asynccontextmanager
|
||
async def _s3_logger_on_production_handler(profile: str, statuses: list[int | type[httpx.TransportError]]):
|
||
"""
|
||
S3Logger wired to the real AsyncHTTPHandler over an httpx MockTransport that answers with the given
|
||
statuses (or raises the given transport errors) in order, so the handler's own error mapping is exercised.
|
||
"""
|
||
requests: list[httpx.Request] = []
|
||
replies = iter(statuses)
|
||
|
||
def respond(request: httpx.Request) -> httpx.Response:
|
||
requests.append(request)
|
||
reply = next(replies)
|
||
if not isinstance(reply, int):
|
||
raise reply("S3 attempt failed", request=request)
|
||
return httpx.Response(reply, request=request, text="<Error><Code>SignatureDoesNotMatch</Code></Error>")
|
||
|
||
handler = AsyncHTTPHandler()
|
||
handler.client = httpx.AsyncClient(transport=httpx.MockTransport(respond))
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_region_name="us-east-1",
|
||
s3_aws_profile_name=profile,
|
||
s3_flush_interval=3600,
|
||
)
|
||
logger.async_httpx_client = handler
|
||
with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep:
|
||
yield logger, requests, mock_sleep
|
||
await handler.client.aclose()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_upload_signs_with_one_frozen_credential_snapshot(rotating_profile: str, caplog):
|
||
"""
|
||
RefreshableCredentials refreshes on every property read once inside the advisory window, so signing
|
||
off the live object would mix the access key of one generation with the token of the next.
|
||
"""
|
||
test_element = s3BatchLoggingElement(
|
||
s3_object_key="2025-09-14/test-frozen.json",
|
||
payload={"test": "frozen"},
|
||
s3_object_download_filename="test-frozen.json",
|
||
)
|
||
async with _s3_logger_on_production_handler(rotating_profile, [200]) as (logger, requests, _):
|
||
await logger.async_upload_data_to_s3(test_element)
|
||
|
||
assert len(requests) == 1
|
||
access_key_generation, token_generation = _generation(requests[0])
|
||
assert access_key_generation == token_generation
|
||
assert "Error uploading to s3" not in caplog.text
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_upload_retries_403_with_fresh_credentials_and_signature(rotating_profile: str, caplog):
|
||
"""
|
||
A 403 (SignatureDoesNotMatch after an IMDS rotation) must be retried, and the retry must fetch
|
||
credentials again and carry a signature computed from that newer generation.
|
||
"""
|
||
test_element = s3BatchLoggingElement(
|
||
s3_object_key="2025-09-14/test-403.json",
|
||
payload={"test": "403"},
|
||
s3_object_download_filename="test-403.json",
|
||
)
|
||
async with _s3_logger_on_production_handler(rotating_profile, [403, 200]) as (logger, requests, mock_sleep):
|
||
await logger.async_upload_data_to_s3(test_element)
|
||
|
||
assert len(requests) == 2
|
||
first_key, first_token = _generation(requests[0])
|
||
second_key, second_token = _generation(requests[1])
|
||
assert first_key == first_token
|
||
assert second_key == second_token
|
||
assert int(second_key) > int(first_key)
|
||
assert requests[1].headers["Authorization"] != requests[0].headers["Authorization"]
|
||
_assert_jittered_backoff(mock_sleep, 1)
|
||
assert "Error uploading to s3" not in caplog.text
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_upload_exhausts_403_retries_through_production_http_handler(rotating_profile: str, caplog):
|
||
test_element = s3BatchLoggingElement(
|
||
s3_object_key="2025-09-14/test-403-exhausted.json",
|
||
payload={"test": "403-exhausted"},
|
||
s3_object_download_filename="test-403-exhausted.json",
|
||
)
|
||
async with _s3_logger_on_production_handler(rotating_profile, [403, 403, 403]) as (logger, requests, mock_sleep):
|
||
await logger.async_upload_data_to_s3(test_element)
|
||
|
||
assert len(requests) == 3
|
||
_assert_jittered_backoff(mock_sleep, 2)
|
||
assert "Error uploading to s3" in caplog.text
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.parametrize("transient_status", [500, 503])
|
||
async def test_async_upload_recovers_from_transient_5xx_through_production_http_handler(
|
||
transient_status: int, rotating_profile: str, caplog: pytest.LogCaptureFixture
|
||
):
|
||
"""
|
||
AsyncHTTPHandler.put raises MaskedHTTPStatusError on 5xx instead of returning the response, so a retry
|
||
loop that only inspects returned status codes never runs (#42868).
|
||
"""
|
||
test_element = s3BatchLoggingElement(
|
||
s3_object_key=f"2025-09-14/test-{transient_status}.json",
|
||
payload={"test": str(transient_status)},
|
||
s3_object_download_filename=f"test-{transient_status}.json",
|
||
)
|
||
async with _s3_logger_on_production_handler(rotating_profile, [transient_status, 200]) as (
|
||
logger,
|
||
requests,
|
||
mock_sleep,
|
||
):
|
||
uploaded = await logger.async_upload_data_to_s3(test_element)
|
||
|
||
assert uploaded is True
|
||
assert len(requests) == 2
|
||
assert all(request.method == "PUT" for request in requests)
|
||
assert requests[0].url == requests[1].url
|
||
assert requests[0].content == requests[1].content
|
||
_assert_jittered_backoff(mock_sleep, 1)
|
||
assert "Error uploading to s3" not in caplog.text
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_upload_is_single_attempted_on_404_through_production_http_handler(rotating_profile: str, caplog):
|
||
test_element = s3BatchLoggingElement(
|
||
s3_object_key="2025-09-14/test-404.json",
|
||
payload={"test": "404"},
|
||
s3_object_download_filename="test-404.json",
|
||
)
|
||
async with _s3_logger_on_production_handler(rotating_profile, [404]) as (logger, requests, mock_sleep):
|
||
outcome = await logger.async_upload_data_to_s3(test_element)
|
||
|
||
assert len(requests) == 1
|
||
assert outcome is False
|
||
mock_sleep.assert_not_awaited()
|
||
assert "Error uploading to s3" in caplog.text
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_upload_access_denied_403_is_retried_and_then_requeued(rotating_profile: str, caplog):
|
||
test_element = s3BatchLoggingElement(
|
||
s3_object_key="2025-09-14/test-403-denied.json",
|
||
payload={"test": "403-denied"},
|
||
s3_object_download_filename="test-403-denied.json",
|
||
)
|
||
requests: list[httpx.Request] = []
|
||
|
||
def respond(request: httpx.Request) -> httpx.Response:
|
||
requests.append(request)
|
||
return httpx.Response(403, request=request, text="<Error><Code>AccessDenied</Code></Error>")
|
||
|
||
handler = AsyncHTTPHandler()
|
||
handler.client = httpx.AsyncClient(transport=httpx.MockTransport(respond))
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_region_name="us-east-1",
|
||
s3_aws_profile_name=rotating_profile,
|
||
s3_flush_interval=3600,
|
||
)
|
||
logger.async_httpx_client = handler
|
||
with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep:
|
||
outcome = await logger.async_upload_data_to_s3(test_element)
|
||
await handler.client.aclose()
|
||
|
||
assert outcome is False
|
||
assert len(requests) == 3
|
||
_assert_jittered_backoff(mock_sleep, 2)
|
||
assert "Error uploading to s3" in caplog.text
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_upload_retries_429_through_production_http_handler(rotating_profile: str, caplog):
|
||
"""S3 and S3-compatible stores answer throttling with 429 too, so it backs off and retries like a 503."""
|
||
async with _s3_logger_on_production_handler(rotating_profile, [429, 200]) as (logger, requests, mock_sleep):
|
||
uploaded = await logger.async_upload_data_to_s3(_element({"test": "429"}, "429"))
|
||
|
||
assert uploaded is True
|
||
assert len(requests) == 2
|
||
assert requests[0].content == requests[1].content
|
||
_assert_jittered_backoff(mock_sleep, 1)
|
||
assert "Error uploading to s3" not in caplog.text
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_upload_retries_a_timed_out_attempt_through_production_http_handler(rotating_profile: str, caplog):
|
||
"""AsyncHTTPHandler.put turns an httpx timeout into litellm.Timeout, which must count as a retryable attempt."""
|
||
async with _s3_logger_on_production_handler(rotating_profile, [httpx.ReadTimeout, 200]) as (
|
||
logger,
|
||
requests,
|
||
mock_sleep,
|
||
):
|
||
uploaded = await logger.async_upload_data_to_s3(_element({"test": "timeout"}, "timeout"))
|
||
|
||
assert uploaded is True
|
||
assert len(requests) == 2
|
||
_assert_jittered_backoff(mock_sleep, 1)
|
||
assert "Error uploading to s3" not in caplog.text
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.parametrize(
|
||
"error",
|
||
[
|
||
pytest.param(httpx.ConnectError("connection refused"), id="connect-error"),
|
||
pytest.param(httpx.RemoteProtocolError("server disconnected"), id="remote-protocol-error"),
|
||
pytest.param(httpx.ReadTimeout("read timed out"), id="read-timeout"),
|
||
],
|
||
)
|
||
async def test_async_upload_retries_transport_errors_then_succeeds(error: Exception) -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = AsyncMock(side_effect=[error, _ok_response()])
|
||
|
||
with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep:
|
||
uploaded = await logger.async_upload_data_to_s3(_element({"test": "transport"}, "transport"))
|
||
|
||
assert uploaded is True
|
||
assert logger.async_httpx_client.put.call_count == 2
|
||
_assert_jittered_backoff(mock_sleep, 1)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_upload_exhausts_transport_error_retries_and_keeps_the_event_queued() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = AsyncMock(side_effect=httpx.ConnectError("connection refused"))
|
||
logger.log_queue = [_element({"test": "down"}, "down")]
|
||
|
||
with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep:
|
||
await logger.flush_queue()
|
||
|
||
assert logger.async_httpx_client.put.call_count == 3
|
||
_assert_jittered_backoff(mock_sleep, 2)
|
||
assert [element.s3_object_key for element in logger.log_queue] == ["2025-09-14/test-down.json"]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_upload_sends_every_attempt_with_a_finite_deadline(rotating_profile: str) -> None:
|
||
"""
|
||
The logger's handler is built with explicit params and no timeout, and the handler forwards that None to
|
||
httpx as "no timeout", so without a per-attempt deadline a stalled S3 endpoint pins the upload forever.
|
||
"""
|
||
async with _s3_logger_on_production_handler(rotating_profile, [503, 200]) as (logger, requests, _):
|
||
await logger.async_upload_data_to_s3(_element({"test": "deadline"}, "deadline"))
|
||
|
||
assert len(requests) == 2
|
||
for request in requests:
|
||
deadlines = request.extensions["timeout"]
|
||
assert all(isinstance(limit, (int, float)) and limit > 0 for limit in deadlines.values()), deadlines
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_backoff_is_jittered_so_throttled_uploads_do_not_retry_in_lockstep() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = AsyncMock(
|
||
side_effect=[_transient_failure_response(), _ok_response(), _transient_failure_response(), _ok_response()]
|
||
)
|
||
|
||
with (
|
||
patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep,
|
||
patch("random.random", side_effect=[0.0, 0.999]),
|
||
):
|
||
await logger.async_upload_data_to_s3(_element({"i": 0}, "0"))
|
||
await logger.async_upload_data_to_s3(_element({"i": 1}, "1"))
|
||
|
||
first, second = (sleep_call.args[0] for sleep_call in mock_sleep.call_args_list)
|
||
assert first != second
|
||
assert 0.5 <= first < 1.5 and 0.5 <= second < 1.5
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_dropped_uploads_are_logged_once_per_flush_naming_each_reason(caplog) -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
terminal = (_element({"i": "terminal-0"}, "terminal-0"), _element({"i": "terminal-1"}, "terminal-1"))
|
||
expired = _element({"i": "expired"}, "expired").model_copy(update={"retrying_since": _NOW - 7200})
|
||
fresh = _element({"i": "fresh"}, "fresh")
|
||
|
||
async def put(
|
||
url: str, data: str | None = None, headers: dict[str, str] | None = None, timeout: httpx.Timeout | None = None
|
||
) -> MagicMock:
|
||
if "test-terminal-" in url:
|
||
return _terminal_failure_response()
|
||
if url.endswith("test-expired.json"):
|
||
return _transient_failure_response()
|
||
return _ok_response()
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
logger.log_queue = [*terminal, expired, fresh]
|
||
|
||
with (
|
||
patch("asyncio.sleep", new_callable=AsyncMock),
|
||
patch("time.monotonic", return_value=_NOW),
|
||
):
|
||
await logger.flush_queue()
|
||
|
||
assert logger.log_queue == []
|
||
drop_lines = [record.getMessage() for record in caplog.records if "uploads dropped" in record.getMessage()]
|
||
assert len(drop_lines) == 1, drop_lines
|
||
assert "3 uploads dropped" in drop_lines[0]
|
||
assert "2 rejected as terminal" in drop_lines[0]
|
||
assert "1 retrying longer than s3_max_retry_age_seconds=3600" in drop_lines[0]
|
||
|
||
|
||
def test_sync_upload_retries_403_with_fresh_signature(rotating_profile: str, monkeypatch: pytest.MonkeyPatch):
|
||
monkeypatch.delenv("AWS_ACCESS_KEY_ID", raising=False)
|
||
monkeypatch.delenv("AWS_SECRET_ACCESS_KEY", raising=False)
|
||
monkeypatch.delenv("AWS_SESSION_TOKEN", raising=False)
|
||
monkeypatch.setenv("AWS_PROFILE", rotating_profile)
|
||
logger = S3Logger(s3_bucket_name="test-bucket", s3_region_name="us-east-1", s3_flush_interval=3600)
|
||
test_element = s3BatchLoggingElement(
|
||
s3_object_key="2025-09-14/test-sync-403.json",
|
||
payload={"test": "sync-403"},
|
||
s3_object_download_filename="test-sync-403.json",
|
||
)
|
||
requests: list[httpx.Request] = []
|
||
replies = iter([403, 200])
|
||
|
||
def respond(request: httpx.Request) -> httpx.Response:
|
||
requests.append(request)
|
||
return httpx.Response(next(replies), request=request, text="<Error><Code>SignatureDoesNotMatch</Code></Error>")
|
||
|
||
handler = HTTPHandler()
|
||
handler.client = httpx.Client(transport=httpx.MockTransport(respond))
|
||
with (
|
||
patch( # test-quality-ok: sync upload builds its HTTPHandler per call, there is no injection seam for it
|
||
"litellm.integrations.s3_v2.get_httpx_client", return_value=handler
|
||
),
|
||
patch("time.sleep") as mock_sleep,
|
||
):
|
||
logger.upload_data_to_s3(test_element)
|
||
|
||
assert len(requests) == 2
|
||
first_key, first_token = _generation(requests[0])
|
||
second_key, second_token = _generation(requests[1])
|
||
assert first_key == first_token
|
||
assert second_key == second_token
|
||
assert int(second_key) > int(first_key)
|
||
_assert_jittered_backoff(mock_sleep, 1)
|
||
|
||
|
||
def test_sync_upload_retries_on_s3_503():
|
||
"""
|
||
Test that the sync upload_data_to_s3 retries on transient S3 503.
|
||
"""
|
||
from unittest.mock import MagicMock
|
||
|
||
from litellm.types.integrations.s3_v2 import s3BatchLoggingElement
|
||
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
|
||
test_element = s3BatchLoggingElement(
|
||
s3_object_key="2025-09-14/test-sync-retry.json",
|
||
payload={"test": "sync-retry"},
|
||
s3_object_download_filename="test-sync-retry.json",
|
||
)
|
||
|
||
response_503 = MagicMock()
|
||
response_503.status_code = 503
|
||
response_200 = MagicMock()
|
||
response_200.status_code = 200
|
||
response_200.raise_for_status = MagicMock()
|
||
|
||
mock_sync_client = MagicMock()
|
||
mock_sync_client.put = MagicMock(side_effect=[response_503, response_200])
|
||
|
||
with patch(
|
||
"litellm.integrations.s3_v2.get_httpx_client",
|
||
return_value=mock_sync_client,
|
||
):
|
||
with patch("time.sleep") as mock_sleep:
|
||
logger.upload_data_to_s3(test_element)
|
||
|
||
assert mock_sync_client.put.call_count == 2
|
||
_assert_jittered_backoff(mock_sleep, 1)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_log_event_skips_when_standard_logging_object_missing():
|
||
"""
|
||
Reproduces the bug where _async_log_event_base raises ValueError when
|
||
kwargs has no standard_logging_object (e.g. call_type=afile_delete).
|
||
|
||
The S3 logger should skip gracefully, not raise.
|
||
"""
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_region_name="us-east-1",
|
||
s3_aws_access_key_id="fake",
|
||
s3_aws_secret_access_key="fake",
|
||
)
|
||
|
||
kwargs_without_slo = {
|
||
"call_type": "afile_delete",
|
||
"model": None,
|
||
"litellm_call_id": "test-call-id",
|
||
}
|
||
|
||
start_time = datetime.utcnow()
|
||
end_time = datetime.utcnow()
|
||
|
||
# Spy on handle_callback_failure — should NOT be called if we skip gracefully.
|
||
# Without the fix, the ValueError is caught by the except block which calls
|
||
# handle_callback_failure. With the fix, we return early and never hit except.
|
||
with patch.object(logger, "handle_callback_failure") as mock_failure:
|
||
await logger._async_log_event_base(
|
||
kwargs=kwargs_without_slo,
|
||
response_obj=None,
|
||
start_time=start_time,
|
||
end_time=end_time,
|
||
)
|
||
|
||
assert not mock_failure.called, (
|
||
"handle_callback_failure should not be called — "
|
||
"missing standard_logging_object should be a graceful skip, not an error"
|
||
)
|
||
|
||
# Nothing should have been queued (catches the case where code falls
|
||
# through without returning and appends None to the queue)
|
||
assert len(logger.log_queue) == 0, "log_queue should be empty when standard_logging_object is missing"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_strip_base64_removes_file_and_nontext_entries():
|
||
logger = S3Logger(s3_strip_base64_files=True)
|
||
|
||
payload = {
|
||
"messages": [
|
||
{
|
||
"role": "user",
|
||
"content": [
|
||
{"type": "text", "text": "Hello world"},
|
||
{
|
||
"type": "image",
|
||
"file": {"file_data": "data:image/png;base64,AAAA"},
|
||
},
|
||
{
|
||
"type": "file",
|
||
"file": {"file_data": "data:application/pdf;base64,BBBB"},
|
||
},
|
||
],
|
||
},
|
||
{
|
||
"role": "assistant",
|
||
"content": [
|
||
{"type": "text", "text": "Response"},
|
||
{
|
||
"type": "audio",
|
||
"file": {"file_data": "data:audio/wav;base64,CCCC"},
|
||
},
|
||
],
|
||
},
|
||
]
|
||
}
|
||
|
||
stripped = await logger._strip_base64_from_messages(payload)
|
||
|
||
# 1️⃣ File/image/audio entries are removed
|
||
assert len(stripped["messages"][0]["content"]) == 1
|
||
assert stripped["messages"][0]["content"][0]["text"] == "Hello world"
|
||
|
||
assert len(stripped["messages"][1]["content"]) == 1
|
||
assert stripped["messages"][1]["content"][0]["text"] == "Response"
|
||
|
||
# 2️⃣ No 'file' keys remain
|
||
for msg in stripped["messages"]:
|
||
for content in msg["content"]:
|
||
assert "file" not in content
|
||
assert content.get("type") == "text"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_strip_base64_keeps_non_file_content():
|
||
logger = S3Logger(s3_strip_base64_files=True)
|
||
|
||
payload = {
|
||
"messages": [
|
||
{
|
||
"role": "user",
|
||
"content": [
|
||
{"type": "text", "text": "Just text"},
|
||
{"type": "text", "text": "Another message"},
|
||
],
|
||
}
|
||
]
|
||
}
|
||
|
||
stripped = await logger._strip_base64_from_messages(payload)
|
||
|
||
# Should not modify pure text messages
|
||
assert stripped["messages"][0]["content"] == payload["messages"][0]["content"]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_strip_base64_handles_empty_or_missing_messages():
|
||
logger = S3Logger(s3_strip_base64_files=True)
|
||
|
||
# Missing messages key
|
||
payload_no_messages = {}
|
||
stripped1 = await logger._strip_base64_from_messages(payload_no_messages)
|
||
assert stripped1 == payload_no_messages
|
||
|
||
# Empty messages list
|
||
payload_empty = {"messages": []}
|
||
stripped2 = await logger._strip_base64_from_messages(payload_empty)
|
||
assert stripped2 == payload_empty
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_strip_base64_mixed_nested_objects():
|
||
"""
|
||
Handles weird/nested content structures gracefully.
|
||
"""
|
||
logger = S3Logger(s3_strip_base64_files=True)
|
||
|
||
payload = {
|
||
"messages": [
|
||
{
|
||
"role": "system",
|
||
"content": [
|
||
{"type": "text", "text": "Keep me"},
|
||
{"type": "custom", "metadata": "ignore but non-text"},
|
||
{"foo": "bar"},
|
||
{"file": {"file_data": "data:application/pdf;base64,XXX"}},
|
||
],
|
||
"extra": {"trace_id": "123"},
|
||
}
|
||
]
|
||
}
|
||
|
||
stripped = await logger._strip_base64_from_messages(payload)
|
||
|
||
# Custom/non-text and file entries removed
|
||
content = stripped["messages"][0]["content"]
|
||
assert len(content) == 2
|
||
assert {"type": "text", "text": "Keep me"} in content
|
||
assert {"foo": "bar"} in content
|
||
# Extra metadata preserved
|
||
assert stripped["messages"][0]["extra"]["trace_id"] == "123"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_s3_verify_false_handling(monkeypatch: pytest.MonkeyPatch):
|
||
"""
|
||
Test that s3_verify=False is properly handled and not treated as None.
|
||
|
||
This is a regression test for the bug where s3_verify=False was being
|
||
ignored because 'False or s3_verify' would evaluate to s3_verify (None).
|
||
"""
|
||
from unittest.mock import AsyncMock, patch
|
||
|
||
import litellm
|
||
|
||
# Set up s3_callback_params with s3_verify=False
|
||
monkeypatch.setattr(
|
||
litellm,
|
||
"s3_callback_params",
|
||
{
|
||
"s3_bucket_name": "test-bucket",
|
||
"s3_endpoint_url": "https://localhost:443",
|
||
"s3_aws_access_key_id": "minioadmin",
|
||
"s3_aws_secret_access_key": "minioadmin",
|
||
"s3_region_name": "us-east-1",
|
||
"s3_verify": False, # This should NOT be ignored
|
||
"s3_use_ssl": False, # This should also NOT be ignored
|
||
},
|
||
)
|
||
|
||
with patch("asyncio.create_task"):
|
||
with patch("litellm.integrations.s3_v2.get_async_httpx_client") as mock_get_client:
|
||
mock_client = AsyncMock()
|
||
mock_get_client.return_value = mock_client
|
||
|
||
# Create logger
|
||
logger = S3Logger()
|
||
|
||
# Verify s3_verify is False, not None
|
||
assert logger.s3_verify is False, f"Expected s3_verify=False, got {logger.s3_verify}"
|
||
assert logger.s3_use_ssl is False, f"Expected s3_use_ssl=False, got {logger.s3_use_ssl}"
|
||
|
||
# Verify that get_async_httpx_client was called with ssl_verify=False
|
||
mock_get_client.assert_called_once()
|
||
call_kwargs = mock_get_client.call_args.kwargs
|
||
assert "params" in call_kwargs, "params should be passed to get_async_httpx_client"
|
||
assert call_kwargs["params"] == {"ssl_verify": False}, (
|
||
f"Expected ssl_verify=False in params, got {call_kwargs.get('params')}"
|
||
)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_s3_verify_none_handling(monkeypatch: pytest.MonkeyPatch):
|
||
"""
|
||
Test that s3_verify=None uses default behavior.
|
||
"""
|
||
from unittest.mock import AsyncMock, patch
|
||
|
||
import litellm
|
||
|
||
# Set up s3_callback_params without s3_verify
|
||
monkeypatch.setattr(
|
||
litellm,
|
||
"s3_callback_params",
|
||
{
|
||
"s3_bucket_name": "test-bucket",
|
||
"s3_aws_access_key_id": "test-key",
|
||
"s3_aws_secret_access_key": "test-secret",
|
||
"s3_region_name": "us-east-1",
|
||
},
|
||
)
|
||
|
||
with patch("asyncio.create_task"):
|
||
with patch("litellm.integrations.s3_v2.get_async_httpx_client") as mock_get_client:
|
||
mock_client = AsyncMock()
|
||
mock_get_client.return_value = mock_client
|
||
|
||
# Create logger without explicit s3_verify
|
||
logger = S3Logger()
|
||
|
||
# Verify s3_verify is None (default)
|
||
assert logger.s3_verify is None, f"Expected s3_verify=None, got {logger.s3_verify}"
|
||
|
||
# Verify that get_async_httpx_client was called
|
||
mock_get_client.assert_called_once()
|
||
call_kwargs = mock_get_client.call_args.kwargs
|
||
# When s3_verify is None, params={'ssl_verify': None} which is fine - uses default behavior
|
||
# The important thing is it's not False
|
||
if "params" in call_kwargs and call_kwargs["params"] is not None:
|
||
assert call_kwargs["params"].get("ssl_verify") is None
|
||
# Either params is None or params={'ssl_verify': None} is acceptable
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_s3_verify_false_creates_httpx_client_with_verify_false(monkeypatch: pytest.MonkeyPatch):
|
||
"""
|
||
Test that when s3_verify=False, the actual httpx client has verify=False.
|
||
|
||
This validates that ssl_verify=False flows through to the httpx.AsyncClient.
|
||
"""
|
||
from unittest.mock import patch
|
||
|
||
import litellm
|
||
|
||
# Set up s3_callback_params with s3_verify=False
|
||
monkeypatch.setattr(
|
||
litellm,
|
||
"s3_callback_params",
|
||
{
|
||
"s3_bucket_name": "test-bucket",
|
||
"s3_endpoint_url": "https://localhost:443",
|
||
"s3_aws_access_key_id": "minioadmin",
|
||
"s3_aws_secret_access_key": "minioadmin",
|
||
"s3_region_name": "us-east-1",
|
||
"s3_verify": False,
|
||
},
|
||
)
|
||
|
||
with patch("asyncio.create_task"):
|
||
# Create logger - this creates the httpx client
|
||
logger = S3Logger()
|
||
|
||
# Verify the logger has s3_verify=False
|
||
assert logger.s3_verify is False
|
||
|
||
# Check the actual httpx client has verify=False
|
||
# The async_httpx_client.client is the actual httpx.AsyncClient
|
||
if hasattr(logger.async_httpx_client, "client"):
|
||
httpx_client = logger.async_httpx_client.client
|
||
# Check the _verify attribute (httpx internal)
|
||
if hasattr(httpx_client, "_verify"):
|
||
assert httpx_client._verify is False, f"Expected httpx client _verify=False, got {httpx_client._verify}"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_s3_verify_false_async_client(monkeypatch: pytest.MonkeyPatch):
|
||
"""
|
||
Test that the async httpx client respects s3_verify=False.
|
||
"""
|
||
from unittest.mock import AsyncMock, MagicMock, patch
|
||
|
||
import litellm
|
||
from litellm.types.integrations.s3_v2 import s3BatchLoggingElement
|
||
|
||
# Set up s3_callback_params with s3_verify=False
|
||
monkeypatch.setattr(
|
||
litellm,
|
||
"s3_callback_params",
|
||
{
|
||
"s3_bucket_name": "test-bucket",
|
||
"s3_endpoint_url": "https://localhost:443",
|
||
"s3_aws_access_key_id": "minioadmin",
|
||
"s3_aws_secret_access_key": "minioadmin",
|
||
"s3_region_name": "us-east-1",
|
||
"s3_verify": False,
|
||
},
|
||
)
|
||
|
||
with patch("asyncio.create_task"):
|
||
logger = S3Logger()
|
||
|
||
# Verify s3_verify is False
|
||
assert logger.s3_verify is False
|
||
|
||
# Create test element
|
||
test_element = s3BatchLoggingElement(
|
||
s3_object_key="2025-11-03/test-key.json",
|
||
payload={"test": "data"},
|
||
s3_object_download_filename="test-file.json",
|
||
)
|
||
|
||
# Mock the async httpx client's put method
|
||
mock_response = MagicMock()
|
||
mock_response.status_code = 200
|
||
mock_response.raise_for_status = MagicMock()
|
||
logger.async_httpx_client.put = AsyncMock(return_value=mock_response)
|
||
|
||
# Call async upload
|
||
await logger.async_upload_data_to_s3(test_element)
|
||
|
||
# Verify put was called
|
||
assert logger.async_httpx_client.put.called
|
||
|
||
# Check that the async httpx client was created with verify=False
|
||
if hasattr(logger.async_httpx_client, "client"):
|
||
httpx_client = logger.async_httpx_client.client
|
||
if hasattr(httpx_client, "_verify"):
|
||
assert httpx_client._verify is False, (
|
||
f"Expected async httpx client _verify=False, got {httpx_client._verify}"
|
||
)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_strip_base64_recursive_redaction():
|
||
logger = S3Logger(s3_strip_base64_files=True)
|
||
payload = {
|
||
"messages": [
|
||
{
|
||
"content": [
|
||
{"type": "text", "text": "normal text"},
|
||
{
|
||
"type": "text",
|
||
"text": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUg",
|
||
},
|
||
{
|
||
"type": "text",
|
||
"text": "Nested: {'data': 'data:application/pdf;base64,AAA...'}",
|
||
},
|
||
{"file": {"file_data": "data:application/pdf;base64,AAAA"}},
|
||
{"metadata": {"preview": "data:audio/mp3;base64,AAAAA=="}},
|
||
]
|
||
}
|
||
]
|
||
}
|
||
|
||
result = await logger._strip_base64_from_messages(payload)
|
||
content = result["messages"][0]["content"]
|
||
|
||
# Dropped file-type entries
|
||
assert not any("file" in c for c in content)
|
||
|
||
# Base64 redacted globally
|
||
import json
|
||
|
||
for c in content:
|
||
if isinstance(c, dict):
|
||
s = json.dumps(c).lower()
|
||
# "[base64_redacted]" is fine, but raw base64 is not
|
||
assert "base64," not in s, f"Found real base64 blob in: {s}"
|
||
|
||
|
||
# --------------------------------------------------------------
|
||
# Shared fixture that silences asyncio.create_task during tests
|
||
# --------------------------------------------------------------
|
||
@pytest_asyncio.fixture(loop_scope="function")
|
||
async def cancel_s3_periodic_flush_tasks(monkeypatch: pytest.MonkeyPatch) -> AsyncIterator[None]:
|
||
periodic_flush_tasks: list[asyncio.Task[object]] = []
|
||
original_create_task = asyncio.create_task
|
||
|
||
def track_create_task(
|
||
coro: Coroutine[object, object, object],
|
||
*,
|
||
name: str | None = None,
|
||
context: Context | None = None,
|
||
) -> asyncio.Task[object]:
|
||
if context is None:
|
||
task = original_create_task(coro, name=name)
|
||
else:
|
||
task = original_create_task(coro, name=name, context=context)
|
||
if coro.__qualname__.endswith("periodic_flush"):
|
||
periodic_flush_tasks.append(task)
|
||
return task
|
||
|
||
monkeypatch.setattr(asyncio, "create_task", track_create_task)
|
||
yield
|
||
for task in periodic_flush_tasks:
|
||
if not task.done():
|
||
task.cancel()
|
||
await asyncio.gather(*periodic_flush_tasks, return_exceptions=True)
|
||
|
||
|
||
@pytest.fixture(autouse=True)
|
||
def patch_asyncio_create_task(request):
|
||
"""Prevent 'no running event loop' errors when S3Logger calls asyncio.create_task()."""
|
||
if request.node.originalname in {
|
||
"test_basic_s3_logging",
|
||
"test_basic_s3_v2_logging",
|
||
"test_basic_s3_v2_logging_failure",
|
||
}:
|
||
yield
|
||
return
|
||
with patch("asyncio.create_task"):
|
||
yield
|
||
|
||
|
||
# --------------------------------------------------------------
|
||
# Parametrized prefix combination test
|
||
# --------------------------------------------------------------
|
||
@pytest.mark.parametrize(
|
||
"use_team_prefix,use_key_prefix,team_alias,key_alias,expected_prefix",
|
||
[
|
||
(False, False, "teamA", "keyA", ""),
|
||
(True, False, "teamA", "keyA", "teamA/"),
|
||
(False, True, "teamA", "keyA", "keyA/"),
|
||
(True, True, "teamA", "keyA", "teamA/keyA/"),
|
||
(True, True, None, "keyA", "keyA/"),
|
||
(True, True, "teamA", None, "teamA/"),
|
||
(True, True, None, None, ""),
|
||
],
|
||
)
|
||
def test_s3_object_key_prefix_combinations(use_team_prefix, use_key_prefix, team_alias, key_alias, expected_prefix):
|
||
"""
|
||
Validate correct S3 prefix composition for team alias + key alias combinations.
|
||
"""
|
||
with patch("litellm.integrations.s3_v2.get_s3_object_key") as mock_get_key:
|
||
mock_get_key.return_value = "mocked/s3/object/key.json"
|
||
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_region_name="us-east-1",
|
||
s3_use_team_prefix=use_team_prefix,
|
||
s3_use_key_prefix=use_key_prefix,
|
||
)
|
||
|
||
payload = StandardLoggingPayload(
|
||
id="abc123",
|
||
metadata={
|
||
"user_api_key_team_alias": team_alias,
|
||
"user_api_key_alias": key_alias,
|
||
},
|
||
messages=[{"role": "user", "content": [{"type": "text", "text": "hi"}]}],
|
||
)
|
||
|
||
result = logger.create_s3_batch_logging_element(datetime.utcnow(), payload)
|
||
assert result is not None
|
||
mock_get_key.assert_called_once()
|
||
|
||
prefix_arg = mock_get_key.call_args.kwargs.get("prefix")
|
||
assert prefix_arg == expected_prefix, (
|
||
f"Expected prefix '{expected_prefix}', got '{prefix_arg}' "
|
||
f"for team={team_alias}, key={key_alias}, "
|
||
f"use_team_prefix={use_team_prefix}, use_key_prefix={use_key_prefix}"
|
||
)
|
||
|
||
|
||
# --------------------------------------------------------------
|
||
# Test prefix priority and concatenation
|
||
# --------------------------------------------------------------
|
||
def test_prefix_priority_and_path_construction():
|
||
"""
|
||
Validate that prefix components are ordered and joined with '/' only once.
|
||
"""
|
||
with patch("litellm.integrations.s3_v2.get_s3_object_key") as mock_get_key:
|
||
mock_get_key.return_value = "mocked/key"
|
||
|
||
logger = S3Logger(s3_use_team_prefix=True, s3_use_key_prefix=True)
|
||
payload = StandardLoggingPayload(
|
||
id="xyz999",
|
||
metadata={
|
||
"user_api_key_team_alias": "Team-Alpha",
|
||
"user_api_key_alias": "API-12345",
|
||
},
|
||
messages=[],
|
||
)
|
||
|
||
logger.create_s3_batch_logging_element(datetime.utcnow(), payload)
|
||
prefix_arg = mock_get_key.call_args.kwargs.get("prefix", "")
|
||
|
||
assert prefix_arg == "Team-Alpha/API-12345/"
|
||
assert "//" not in prefix_arg
|
||
|
||
|
||
# --------------------------------------------------------------
|
||
# Test when prefixes are disabled
|
||
# --------------------------------------------------------------
|
||
def test_prefix_absent_when_flags_disabled():
|
||
"""
|
||
Verify prefix is omitted entirely when prefix flags are False.
|
||
"""
|
||
with patch("litellm.integrations.s3_v2.get_s3_object_key") as mock_get_key:
|
||
mock_get_key.return_value = "mocked/key"
|
||
|
||
logger = S3Logger(s3_use_team_prefix=False, s3_use_key_prefix=False)
|
||
payload = StandardLoggingPayload(
|
||
id="no-prefix",
|
||
metadata={
|
||
"user_api_key_team_alias": "team-x",
|
||
"user_api_key_alias": "key-x",
|
||
},
|
||
messages=[],
|
||
)
|
||
|
||
logger.create_s3_batch_logging_element(datetime.utcnow(), payload)
|
||
prefix_arg = mock_get_key.call_args.kwargs.get("prefix", None)
|
||
assert prefix_arg == "", f"Expected empty prefix, got {prefix_arg}"
|
||
|
||
|
||
# --------------------------------------------------------------
|
||
# Integration-style test (asyncio fixture will patch create_task)
|
||
# --------------------------------------------------------------
|
||
@pytest.mark.asyncio
|
||
async def test_combined_prefix_reflects_in_s3_object_key():
|
||
"""
|
||
Integration-style test ensuring final s3_object_key includes both prefixes correctly.
|
||
"""
|
||
logger = S3Logger(s3_use_team_prefix=True, s3_use_key_prefix=True)
|
||
payload = StandardLoggingPayload(
|
||
id="int-test",
|
||
metadata={
|
||
"user_api_key_team_alias": "myteam",
|
||
"user_api_key_alias": "apikey",
|
||
},
|
||
messages=[],
|
||
)
|
||
|
||
result = logger.create_s3_batch_logging_element(datetime.utcnow(), payload)
|
||
key = result.s3_object_key
|
||
assert "myteam/apikey/" in key, f"Expected both prefixes in key: {key}"
|
||
|
||
|
||
def test_s3_object_key_sanitizes_slashes_and_colons_in_file_name():
|
||
from litellm.integrations.s3 import get_s3_object_key
|
||
|
||
start_time = datetime(2026, 2, 11, 0, 35, 18, 391582)
|
||
file_name = "time-00-35-18-391582_arn:aws:bedrock:us-east-1:123456789012:model-invocation-job/gl18r6skk9yy"
|
||
|
||
key = get_s3_object_key(
|
||
s3_path="LiteLLMAPPLogs",
|
||
prefix="myteam/",
|
||
start_time=start_time,
|
||
s3_file_name=file_name,
|
||
)
|
||
|
||
assert key == (
|
||
"LiteLLMAPPLogs/myteam/2026-02-11/"
|
||
"time-00-35-18-391582_arn_aws_bedrock_us-east-1_123456789012_model-invocation-job_gl18r6skk9yy.json"
|
||
)
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
"response_id",
|
||
[
|
||
"s3://example-batch-bucket/litellm-bedrock-files/input.jsonl",
|
||
"gs://example-batch-bucket/litellm-vertex-files/input.jsonl",
|
||
],
|
||
)
|
||
def test_s3_object_key_has_no_colon_for_cloud_uri_file_ids(response_id: str):
|
||
from litellm.integrations.s3 import get_s3_object_key
|
||
|
||
key = get_s3_object_key(
|
||
s3_path="",
|
||
prefix="",
|
||
start_time=datetime(2026, 9, 7, 4, 51, 6, 685889),
|
||
s3_file_name=f"time-04-51-06-685889_{response_id}",
|
||
)
|
||
|
||
filename = key.rsplit("/", 1)[-1]
|
||
assert ":" not in filename
|
||
assert filename.endswith("_input.jsonl.json")
|
||
|
||
|
||
def test_create_s3_batch_logging_element_flat_key_for_arn_response_id():
|
||
"""End-to-end through the s3_v2 element builder: an ARN response id must
|
||
yield a flat file directly under the date segment."""
|
||
logger = S3Logger(s3_use_team_prefix=False, s3_use_key_prefix=False)
|
||
payload = StandardLoggingPayload(
|
||
id="arn:aws:bedrock:us-east-1:123456789012:model-invocation-job/gl18r6skk9yy",
|
||
metadata={},
|
||
messages=[],
|
||
)
|
||
|
||
start_time = datetime(2026, 2, 11, 0, 35, 18, 391582)
|
||
result = logger.create_s3_batch_logging_element(start_time, payload)
|
||
|
||
assert result is not None
|
||
date_segment = "2026-02-11/"
|
||
file_segment = result.s3_object_key.split(date_segment, 1)[1]
|
||
assert "/" not in file_segment, f"Expected flat file under date segment, got: {result.s3_object_key}"
|
||
assert file_segment.endswith("model-invocation-job_gl18r6skk9yy.json")
|
||
|
||
|
||
# --------------------------------------------------------------
|
||
# object keys bounded to S3's 1024 UTF-8 byte limit
|
||
# --------------------------------------------------------------
|
||
def _oversized_response_id() -> str:
|
||
return "resp_" + "A" * 1100
|
||
|
||
|
||
def test_s3_object_key_at_the_byte_limit_is_left_alone():
|
||
"""A key that still fits is left byte-identical."""
|
||
from litellm.constants import MAX_S3_OBJECT_KEY_BYTES
|
||
from litellm.integrations.s3 import get_s3_object_key
|
||
|
||
start_time = datetime(2026, 8, 24, 6, 18, 41, 948021)
|
||
fixed_len = len("input/2026-08-24/.json")
|
||
file_name = "x" * (MAX_S3_OBJECT_KEY_BYTES - fixed_len)
|
||
|
||
key = get_s3_object_key(s3_path="input", prefix="", start_time=start_time, s3_file_name=file_name)
|
||
|
||
assert key == f"input/2026-08-24/{file_name}.json"
|
||
assert len(key.encode("utf-8")) == MAX_S3_OBJECT_KEY_BYTES
|
||
|
||
|
||
def test_s3_object_key_is_bounded_for_oversized_response_id():
|
||
"""An oversized Responses API id is shortened to a readable head plus a digest."""
|
||
import hashlib
|
||
|
||
from litellm.constants import MAX_S3_OBJECT_KEY_BYTES
|
||
from litellm.integrations.s3 import get_s3_object_key
|
||
|
||
start_time = datetime(2026, 8, 24, 6, 18, 41, 948021)
|
||
file_name = f"time-06-18-41-948021_{_oversized_response_id()}"
|
||
|
||
key = get_s3_object_key(s3_path="input", prefix="DefaultTeamProd/", start_time=start_time, s3_file_name=file_name)
|
||
|
||
assert len(key.encode("utf-8")) <= MAX_S3_OBJECT_KEY_BYTES
|
||
assert key.startswith("input/DefaultTeamProd/2026-08-24/time-06-18-41-948021_resp_")
|
||
assert key.endswith(f"_{hashlib.sha256(file_name.encode('utf-8')).hexdigest()}.json")
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
"s3_path,prefix",
|
||
[
|
||
("input", ""),
|
||
("a" * 900, ""),
|
||
("input", "team-" + "b" * 900 + "/"),
|
||
("c" * 600, "team-" + "d" * 600 + "/key-" + "e" * 600 + "/"),
|
||
# many short segments, so the trim lands exactly on the budget edge
|
||
("", "ssss/" * 200),
|
||
],
|
||
)
|
||
def test_s3_object_key_is_bounded_for_long_paths_and_aliases(s3_path: str, prefix: str):
|
||
"""Long paths, team aliases and key aliases stay within the cap."""
|
||
from litellm.constants import MAX_S3_OBJECT_KEY_BYTES
|
||
from litellm.integrations.s3 import get_s3_object_key
|
||
|
||
key = get_s3_object_key(
|
||
s3_path=s3_path,
|
||
prefix=prefix,
|
||
start_time=datetime(2026, 8, 24, 6, 18, 41, 948021),
|
||
s3_file_name=f"time-06-18-41-948021_{_oversized_response_id()}",
|
||
)
|
||
|
||
assert len(key.encode("utf-8")) <= MAX_S3_OBJECT_KEY_BYTES
|
||
assert key.endswith(".json")
|
||
assert "/2026-08-24/" in key or key.startswith("2026-08-24/")
|
||
assert "/" not in key.rsplit("2026-08-24/", 1)[1]
|
||
|
||
|
||
def test_s3_object_key_trimmed_prefixes_stay_distinct_per_operator():
|
||
"""Prefixes that differ only past the trim point keep separate folders."""
|
||
from litellm.constants import MAX_S3_OBJECT_KEY_BYTES
|
||
from litellm.integrations.s3 import get_s3_object_key
|
||
|
||
start_time = datetime(2026, 8, 24, 6, 18, 41, 948021)
|
||
keys = [
|
||
get_s3_object_key(
|
||
s3_path="input",
|
||
prefix="team-" + "b" * 1000 + suffix + "/",
|
||
start_time=start_time,
|
||
s3_file_name=f"time-06-18-41-948021_{_oversized_response_id()}",
|
||
)
|
||
for suffix in ("-one", "-two")
|
||
]
|
||
|
||
assert keys[0] != keys[1]
|
||
assert all(key.startswith("input/team-" + "b" * 900) for key in keys)
|
||
assert all(len(key.encode("utf-8")) == MAX_S3_OBJECT_KEY_BYTES for key in keys)
|
||
|
||
|
||
def test_s3_object_key_bounded_prefix_never_splits_a_multibyte_character():
|
||
"""A multibyte prefix is trimmed on a character boundary."""
|
||
from litellm.constants import MAX_S3_OBJECT_KEY_BYTES
|
||
from litellm.integrations.s3 import get_s3_object_key
|
||
|
||
s3_path = "\u65e5\u672c\u8a9e" * 200
|
||
|
||
key = get_s3_object_key(
|
||
s3_path=s3_path,
|
||
prefix="\u30c1\u30fc\u30e0" * 200 + "/",
|
||
start_time=datetime(2026, 8, 24, 6, 18, 41, 948021),
|
||
s3_file_name=f"time-06-18-41-948021_{_oversized_response_id()}",
|
||
)
|
||
|
||
assert len(key.encode("utf-8")) <= MAX_S3_OBJECT_KEY_BYTES
|
||
assert key.startswith(s3_path[:100])
|
||
assert "\ufffd" not in key
|
||
|
||
|
||
def test_s3_object_key_stays_unique_for_ids_sharing_a_head():
|
||
"""Ids sharing a visible head still get distinct keys."""
|
||
from litellm.integrations.s3 import get_s3_object_key
|
||
|
||
start_time = datetime(2026, 8, 24, 6, 18, 41, 948021)
|
||
keys = {
|
||
get_s3_object_key(
|
||
s3_path="input",
|
||
prefix="",
|
||
start_time=start_time,
|
||
s3_file_name=f"time-06-18-41-948021_{_oversized_response_id()}{suffix}",
|
||
)
|
||
for suffix in ("first", "second", "third")
|
||
}
|
||
|
||
assert len(keys) == 3
|
||
|
||
|
||
def test_s3_object_key_bounding_matches_the_documented_layout():
|
||
"""The bounded key is `<prefix>/<date>/<head>_<sha256>.json`."""
|
||
import hashlib
|
||
|
||
from litellm.integrations.s3 import get_s3_object_key
|
||
|
||
file_name = f"time-06-18-41-948021_{_oversized_response_id()}"
|
||
|
||
key = get_s3_object_key(
|
||
s3_path="input",
|
||
prefix="team/",
|
||
start_time=datetime(2026, 8, 24, 6, 18, 41, 948021),
|
||
s3_file_name=file_name,
|
||
)
|
||
|
||
digest = hashlib.sha256(file_name.encode("utf-8")).hexdigest()
|
||
assert key == f"input/team/2026-08-24/{file_name[:64]}_{digest}.json"
|
||
|
||
|
||
def test_s3_object_key_keeps_the_configured_prefix_when_only_the_id_overflows():
|
||
"""A 940 byte configured prefix survives whole when only the id overflows."""
|
||
from litellm.constants import MAX_S3_OBJECT_KEY_BYTES
|
||
from litellm.integrations.s3 import get_s3_object_key
|
||
|
||
prefix = "team-" + "b" * 934 + "/"
|
||
|
||
key = get_s3_object_key(
|
||
s3_path="",
|
||
prefix=prefix,
|
||
start_time=datetime(2026, 8, 24, 6, 18, 41, 948021),
|
||
s3_file_name=f"time-06-18-41-948021_{_oversized_response_id()}",
|
||
)
|
||
|
||
assert key.startswith(prefix + "2026-08-24/")
|
||
assert len(key.encode("utf-8")) == MAX_S3_OBJECT_KEY_BYTES
|
||
|
||
|
||
def test_s3_object_key_spends_the_whole_budget_when_the_prefix_must_be_trimmed():
|
||
"""A trimmed prefix keeps every byte the budget allows, not whole segments."""
|
||
from litellm.constants import MAX_S3_OBJECT_KEY_BYTES
|
||
from litellm.integrations.s3 import get_s3_object_key
|
||
|
||
s3_path = "p" * 400 + "/" + "q" * 600
|
||
|
||
key = get_s3_object_key(
|
||
s3_path=s3_path,
|
||
prefix="",
|
||
start_time=datetime(2026, 8, 24, 6, 18, 41, 948021),
|
||
s3_file_name="time-06-18-41-948021_abc",
|
||
)
|
||
|
||
assert len(key.encode("utf-8")) == MAX_S3_OBJECT_KEY_BYTES
|
||
assert key.startswith("p" * 400 + "/" + "q" * 500)
|
||
|
||
|
||
def test_s3_object_key_keeps_a_single_segment_path_as_far_as_it_fits():
|
||
"""A path with no separator is kept as far as it fits, never dropped to the bucket root."""
|
||
from litellm.constants import MAX_S3_OBJECT_KEY_BYTES
|
||
from litellm.integrations.s3 import get_s3_object_key
|
||
|
||
key = get_s3_object_key(
|
||
s3_path="a" * 1050,
|
||
prefix="",
|
||
start_time=datetime(2026, 8, 24, 6, 18, 41, 948021),
|
||
s3_file_name="time-06-18-41-948021_chatcmpl-xyz",
|
||
)
|
||
|
||
assert len(key.encode("utf-8")) == MAX_S3_OBJECT_KEY_BYTES
|
||
assert key.startswith("a" * 900)
|
||
|
||
|
||
def test_create_s3_batch_logging_element_bounds_key_and_keeps_full_response_id():
|
||
"""The batch element bounds the key and keeps the full response id in the payload."""
|
||
from litellm.constants import MAX_S3_OBJECT_KEY_BYTES
|
||
|
||
logger = S3Logger(s3_use_team_prefix=True, s3_use_key_prefix=True)
|
||
response_id = _oversized_response_id()
|
||
payload = StandardLoggingPayload(
|
||
id=response_id,
|
||
metadata={"user_api_key_team_alias": "DefaultTeamProd", "user_api_key_alias": "prod-key"},
|
||
messages=[],
|
||
)
|
||
|
||
result = logger.create_s3_batch_logging_element(datetime(2026, 8, 24, 6, 18, 41, 948021), payload)
|
||
|
||
assert result is not None
|
||
assert len(result.s3_object_key.encode("utf-8")) <= MAX_S3_OBJECT_KEY_BYTES
|
||
assert result.s3_object_key.startswith("DefaultTeamProd/prod-key/2026-08-24/")
|
||
assert result.payload["id"] == response_id
|
||
|
||
|
||
def test_s3_object_download_filename_is_bounded_for_oversized_response_id():
|
||
"""The Content-Disposition filename is bounded too, or the PUT fails with MetadataTooLarge."""
|
||
from litellm.constants import MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES
|
||
from litellm.integrations.s3 import get_s3_object_download_filename
|
||
|
||
file_name = get_s3_object_download_filename(datetime(2026, 8, 24, 6, 18, 41, 948021), _oversized_response_id())
|
||
|
||
assert len(file_name.encode("utf-8")) <= MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES
|
||
assert file_name.startswith("time-2026-08-24T06-18-41-948021_resp_")
|
||
assert file_name.endswith(".json")
|
||
|
||
|
||
def test_s3_object_download_filenames_stay_distinct_when_shortened():
|
||
"""Shortened filenames stay distinct."""
|
||
from litellm.integrations.s3 import get_s3_object_download_filename
|
||
|
||
start_time = datetime(2026, 8, 24, 6, 18, 41, 948021)
|
||
file_names = {
|
||
get_s3_object_download_filename(start_time, _oversized_response_id() + suffix)
|
||
for suffix in ("first", "second", "third")
|
||
}
|
||
|
||
assert len(file_names) == 3
|
||
|
||
|
||
def test_s3_object_download_filename_short_id_is_unchanged():
|
||
"""An ordinary response id keeps the filename it had before."""
|
||
from litellm.integrations.s3 import get_s3_object_download_filename
|
||
|
||
file_name = get_s3_object_download_filename(datetime(2026, 8, 24, 6, 18, 41, 948021), "resp_abc123")
|
||
|
||
assert file_name == "time-2026-08-24T06-18-41-948021_resp_abc123.json"
|
||
|
||
|
||
def test_create_s3_batch_logging_element_bounds_the_download_filename():
|
||
"""The batch element carries a bounded Content-Disposition filename."""
|
||
from litellm.constants import MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES
|
||
|
||
logger = S3Logger()
|
||
payload = StandardLoggingPayload(id=_oversized_response_id(), metadata={}, messages=[])
|
||
|
||
result = logger.create_s3_batch_logging_element(datetime(2026, 8, 24, 6, 18, 41, 948021), payload)
|
||
|
||
assert result is not None
|
||
assert len(result.s3_object_download_filename.encode("utf-8")) <= MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_audit_log_object_key_is_bounded_for_a_long_configured_path():
|
||
"""Audit log keys are bounded by the same builder."""
|
||
from litellm.constants import MAX_S3_OBJECT_KEY_BYTES
|
||
|
||
logger = S3Logger()
|
||
logger.s3_path = "audit-archive/" + "z" * 1100
|
||
|
||
await logger.async_log_audit_log_event({"id": "1a4f7bd0-6f1e-4d0a-9b3c-9f2e1d5a7c88"})
|
||
|
||
assert len(logger.log_queue) == 1
|
||
assert len(logger.log_queue[0].s3_object_key.encode("utf-8")) <= MAX_S3_OBJECT_KEY_BYTES
|
||
assert logger.log_queue[0].s3_object_key.startswith("audit-archive/" + "z" * 900)
|
||
|
||
|
||
def test_s3_object_download_filename_drops_characters_that_break_the_header():
|
||
"""A quote or separator in the response id cannot escape the quoted header value."""
|
||
from litellm.integrations.s3 import get_s3_object_download_filename
|
||
|
||
file_name = get_s3_object_download_filename(datetime(2026, 8, 24, 6, 18, 41, 948021), 'resp_a"b/c')
|
||
|
||
assert file_name == "time-2026-08-24T06-18-41-948021_resp_a_b_c.json"
|
||
|
||
|
||
# --------------------------------------------------------------
|
||
# params_source / s3_callback_params_override (audit-log decoupling)
|
||
# --------------------------------------------------------------
|
||
def test_s3_callback_params_override_uses_alternate_dict(monkeypatch):
|
||
"""`s3_callback_params_override` makes the logger read its config from
|
||
the override dict instead of `litellm.s3_callback_params`."""
|
||
import litellm
|
||
|
||
monkeypatch.setattr(litellm, "s3_callback_params", {"s3_bucket_name": "normal-bucket"})
|
||
logger = S3Logger(
|
||
s3_callback_params_override={
|
||
"s3_bucket_name": "audit-bucket",
|
||
"s3_path": "audit-prefix",
|
||
"s3_region_name": "us-west-2",
|
||
}
|
||
)
|
||
assert logger.s3_bucket_name == "audit-bucket"
|
||
assert logger.s3_path == "audit-prefix"
|
||
assert logger.s3_region_name == "us-west-2"
|
||
|
||
|
||
def test_s3_callback_params_override_does_not_mutate_inputs(monkeypatch):
|
||
"""Resolving `os.environ/X` markers must not mutate the override dict
|
||
or `litellm.s3_callback_params`."""
|
||
import litellm
|
||
|
||
monkeypatch.setenv("MY_AUDIT_BUCKET", "resolved-bucket")
|
||
override = {"s3_bucket_name": "os.environ/MY_AUDIT_BUCKET"}
|
||
monkeypatch.setattr(litellm, "s3_callback_params", {"s3_bucket_name": "os.environ/MY_AUDIT_BUCKET"})
|
||
logger = S3Logger(s3_callback_params_override=override)
|
||
assert logger.s3_bucket_name == "resolved-bucket"
|
||
assert override["s3_bucket_name"] == "os.environ/MY_AUDIT_BUCKET"
|
||
assert litellm.s3_callback_params["s3_bucket_name"] == "os.environ/MY_AUDIT_BUCKET"
|
||
|
||
|
||
def test_s3_callback_params_override_none_falls_back_to_global(monkeypatch):
|
||
"""No override → behaves exactly as today (reads `litellm.s3_callback_params`)."""
|
||
import litellm
|
||
|
||
monkeypatch.setattr(litellm, "s3_callback_params", {"s3_bucket_name": "from-global"})
|
||
logger = S3Logger()
|
||
assert logger.s3_bucket_name == "from-global"
|
||
|
||
|
||
def test_s3_callback_params_override_empty_dict_is_opt_in(monkeypatch):
|
||
"""An empty override dict skips the global entirely (env/IAM-only config)."""
|
||
import litellm
|
||
|
||
monkeypatch.setattr(litellm, "s3_callback_params", {"s3_bucket_name": "from-global"})
|
||
logger = S3Logger(s3_callback_params_override={})
|
||
assert logger.s3_bucket_name is None
|
||
|
||
|
||
def _expected_content_md5(payload: dict) -> str:
|
||
import base64
|
||
import hashlib
|
||
|
||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||
|
||
json_string = safe_dumps(payload)
|
||
return base64.b64encode(hashlib.md5(json_string.encode("utf-8"), usedforsecurity=False).digest()).decode()
|
||
|
||
|
||
def _require_non_security_md5(monkeypatch):
|
||
import hashlib
|
||
|
||
original_md5 = hashlib.md5
|
||
|
||
def fips_md5(data=b"", *, usedforsecurity=True):
|
||
if usedforsecurity:
|
||
raise ValueError("MD5 blocked for security use")
|
||
return original_md5(data, usedforsecurity=usedforsecurity)
|
||
|
||
monkeypatch.setattr(hashlib, "md5", fips_md5)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_upload_sets_content_md5_header(monkeypatch):
|
||
"""
|
||
Object Lock buckets reject PUTs without a Content-MD5 header (AWS spec).
|
||
The async upload must send a base64 md5 of the exact signed body.
|
||
"""
|
||
from unittest.mock import AsyncMock, MagicMock
|
||
|
||
from litellm.types.integrations.s3_v2 import s3BatchLoggingElement
|
||
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
|
||
payload = {"test": "content-md5"}
|
||
test_element = s3BatchLoggingElement(
|
||
s3_object_key="2025-09-14/test-md5.json",
|
||
payload=payload,
|
||
s3_object_download_filename="test-md5.json",
|
||
)
|
||
_require_non_security_md5(monkeypatch)
|
||
|
||
response = MagicMock()
|
||
response.status_code = 200
|
||
response.raise_for_status = MagicMock()
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put.return_value = response
|
||
|
||
await logger.async_upload_data_to_s3(test_element)
|
||
|
||
headers = logger.async_httpx_client.put.call_args.kwargs["headers"]
|
||
assert headers["Content-MD5"] == _expected_content_md5(payload)
|
||
assert "x-amz-server-side-encryption" not in headers
|
||
|
||
|
||
def test_sync_upload_sets_content_md5_header(monkeypatch):
|
||
"""The sync upload path must also send Content-MD5 for Object Lock buckets."""
|
||
from unittest.mock import MagicMock
|
||
|
||
from litellm.types.integrations.s3_v2 import s3BatchLoggingElement
|
||
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
|
||
payload = {"test": "sync-content-md5"}
|
||
test_element = s3BatchLoggingElement(
|
||
s3_object_key="2025-09-14/test-sync-md5.json",
|
||
payload=payload,
|
||
s3_object_download_filename="test-sync-md5.json",
|
||
)
|
||
_require_non_security_md5(monkeypatch)
|
||
|
||
response = MagicMock()
|
||
response.status_code = 200
|
||
response.raise_for_status = MagicMock()
|
||
mock_sync_client = MagicMock()
|
||
mock_sync_client.put.return_value = response
|
||
|
||
with patch(
|
||
"litellm.integrations.s3_v2.get_httpx_client",
|
||
return_value=mock_sync_client,
|
||
):
|
||
logger.upload_data_to_s3(test_element)
|
||
|
||
headers = mock_sync_client.put.call_args.kwargs["headers"]
|
||
assert headers["Content-MD5"] == _expected_content_md5(payload)
|
||
assert "x-amz-server-side-encryption" not in headers
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_upload_sets_server_side_encryption_header_when_configured():
|
||
"""
|
||
When s3_server_side_encryption is set (e.g. buckets with a KMS default
|
||
encryption policy), the PUT must carry x-amz-server-side-encryption.
|
||
"""
|
||
from unittest.mock import AsyncMock, MagicMock
|
||
|
||
from litellm.types.integrations.s3_v2 import s3BatchLoggingElement
|
||
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_server_side_encryption="aws:kms",
|
||
)
|
||
|
||
test_element = s3BatchLoggingElement(
|
||
s3_object_key="2025-09-14/test-sse.json",
|
||
payload={"test": "sse"},
|
||
s3_object_download_filename="test-sse.json",
|
||
)
|
||
|
||
response = MagicMock()
|
||
response.status_code = 200
|
||
response.raise_for_status = MagicMock()
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put.return_value = response
|
||
|
||
await logger.async_upload_data_to_s3(test_element)
|
||
|
||
headers = logger.async_httpx_client.put.call_args.kwargs["headers"]
|
||
assert headers["x-amz-server-side-encryption"] == "aws:kms"
|
||
|
||
|
||
def test_s3_server_side_encryption_read_from_callback_params(monkeypatch):
|
||
"""s3_server_side_encryption can be configured via s3_callback_params."""
|
||
import litellm
|
||
|
||
monkeypatch.setattr(
|
||
litellm,
|
||
"s3_callback_params",
|
||
{
|
||
"s3_bucket_name": "from-global",
|
||
"s3_server_side_encryption": "aws:kms",
|
||
},
|
||
)
|
||
logger = S3Logger()
|
||
assert logger.s3_server_side_encryption == "aws:kms"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_upload_sets_sse_kms_key_id_header_when_configured():
|
||
"""
|
||
When s3_sse_kms_key_id is set alongside aws:kms, the PUT must carry
|
||
x-amz-server-side-encryption-aws-kms-key-id so objects are encrypted
|
||
with the customer-managed KMS key instead of the bucket default.
|
||
"""
|
||
from unittest.mock import AsyncMock, MagicMock
|
||
|
||
from litellm.types.integrations.s3_v2 import s3BatchLoggingElement
|
||
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_server_side_encryption="aws:kms",
|
||
s3_sse_kms_key_id="arn:aws:kms:us-east-1:111122223333:key/test-key-id",
|
||
)
|
||
|
||
test_element = s3BatchLoggingElement(
|
||
s3_object_key="2025-09-14/test-sse-kms.json",
|
||
payload={"test": "sse-kms"},
|
||
s3_object_download_filename="test-sse-kms.json",
|
||
)
|
||
|
||
response = MagicMock()
|
||
response.status_code = 200
|
||
response.raise_for_status = MagicMock()
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put.return_value = response
|
||
|
||
await logger.async_upload_data_to_s3(test_element)
|
||
|
||
headers = logger.async_httpx_client.put.call_args.kwargs["headers"]
|
||
assert headers["x-amz-server-side-encryption"] == "aws:kms"
|
||
assert headers["x-amz-server-side-encryption-aws-kms-key-id"] == (
|
||
"arn:aws:kms:us-east-1:111122223333:key/test-key-id"
|
||
)
|
||
|
||
|
||
def test_sync_upload_sets_sse_kms_key_id_header_when_configured():
|
||
"""The sync upload path must carry the same SSE-KMS headers."""
|
||
from unittest.mock import MagicMock
|
||
|
||
from litellm.types.integrations.s3_v2 import s3BatchLoggingElement
|
||
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_server_side_encryption="aws:kms",
|
||
s3_sse_kms_key_id="arn:aws:kms:us-east-1:111122223333:key/test-key-id",
|
||
)
|
||
|
||
test_element = s3BatchLoggingElement(
|
||
s3_object_key="2025-09-14/test-sync-sse-kms.json",
|
||
payload={"test": "sync-sse-kms"},
|
||
s3_object_download_filename="test-sync-sse-kms.json",
|
||
)
|
||
|
||
response = MagicMock()
|
||
response.status_code = 200
|
||
response.raise_for_status = MagicMock()
|
||
mock_sync_client = MagicMock()
|
||
mock_sync_client.put.return_value = response
|
||
|
||
with patch(
|
||
"litellm.integrations.s3_v2.get_httpx_client",
|
||
return_value=mock_sync_client,
|
||
):
|
||
logger.upload_data_to_s3(test_element)
|
||
|
||
headers = mock_sync_client.put.call_args.kwargs["headers"]
|
||
assert headers["x-amz-server-side-encryption"] == "aws:kms"
|
||
assert headers["x-amz-server-side-encryption-aws-kms-key-id"] == (
|
||
"arn:aws:kms:us-east-1:111122223333:key/test-key-id"
|
||
)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_upload_omits_kms_key_id_header_when_not_configured():
|
||
"""SSE without a key id must not emit the KMS key id header."""
|
||
from unittest.mock import AsyncMock, MagicMock
|
||
|
||
from litellm.types.integrations.s3_v2 import s3BatchLoggingElement
|
||
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_server_side_encryption="AES256",
|
||
)
|
||
|
||
test_element = s3BatchLoggingElement(
|
||
s3_object_key="2025-09-14/test-aes256.json",
|
||
payload={"test": "aes256"},
|
||
s3_object_download_filename="test-aes256.json",
|
||
)
|
||
|
||
response = MagicMock()
|
||
response.status_code = 200
|
||
response.raise_for_status = MagicMock()
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put.return_value = response
|
||
|
||
await logger.async_upload_data_to_s3(test_element)
|
||
|
||
headers = logger.async_httpx_client.put.call_args.kwargs["headers"]
|
||
assert headers["x-amz-server-side-encryption"] == "AES256"
|
||
assert "x-amz-server-side-encryption-aws-kms-key-id" not in headers
|
||
|
||
|
||
def test_s3_sse_kms_key_id_read_from_callback_params(monkeypatch):
|
||
"""s3_sse_kms_key_id can be configured via s3_callback_params."""
|
||
import litellm
|
||
|
||
monkeypatch.setattr(
|
||
litellm,
|
||
"s3_callback_params",
|
||
{
|
||
"s3_bucket_name": "from-global",
|
||
"s3_server_side_encryption": "aws:kms",
|
||
"s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id",
|
||
},
|
||
)
|
||
logger = S3Logger()
|
||
assert logger.s3_sse_kms_key_id == ("arn:aws:kms:us-east-1:111122223333:key/test-key-id")
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_upload_infers_aws_kms_when_only_key_id_set():
|
||
"""
|
||
Setting only s3_sse_kms_key_id must not produce an invalid request
|
||
(S3 rejects a key id without an algorithm); aws:kms is inferred.
|
||
"""
|
||
from unittest.mock import AsyncMock, MagicMock
|
||
|
||
from litellm.types.integrations.s3_v2 import s3BatchLoggingElement
|
||
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_sse_kms_key_id="arn:aws:kms:us-east-1:111122223333:key/test-key-id",
|
||
)
|
||
|
||
test_element = s3BatchLoggingElement(
|
||
s3_object_key="2025-09-14/test-kms-only.json",
|
||
payload={"test": "kms-only"},
|
||
s3_object_download_filename="test-kms-only.json",
|
||
)
|
||
|
||
response = MagicMock()
|
||
response.status_code = 200
|
||
response.raise_for_status = MagicMock()
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put.return_value = response
|
||
|
||
await logger.async_upload_data_to_s3(test_element)
|
||
|
||
headers = logger.async_httpx_client.put.call_args.kwargs["headers"]
|
||
assert headers["x-amz-server-side-encryption"] == "aws:kms"
|
||
assert headers["x-amz-server-side-encryption-aws-kms-key-id"] == (
|
||
"arn:aws:kms:us-east-1:111122223333:key/test-key-id"
|
||
)
|
||
|
||
|
||
def test_s3_sse_kms_key_id_read_from_audit_override_params(monkeypatch):
|
||
"""The audit-log override path must honor s3_sse_kms_key_id too."""
|
||
import litellm
|
||
|
||
monkeypatch.setattr(litellm, "s3_callback_params", {"s3_bucket_name": "normal-logs-bucket"})
|
||
logger = S3Logger(
|
||
s3_callback_params_override={
|
||
"s3_bucket_name": "audit-logs-bucket",
|
||
"s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/audit-key-id",
|
||
}
|
||
)
|
||
assert logger.s3_bucket_name == "audit-logs-bucket"
|
||
assert logger.s3_sse_kms_key_id == ("arn:aws:kms:us-east-1:111122223333:key/audit-key-id")
|
||
|
||
|
||
def test_kms_key_id_dropped_when_algorithm_is_not_kms(monkeypatch):
|
||
"""
|
||
AES256 plus a KMS key id is an invalid S3 combination; the key id must be
|
||
dropped at init so uploads keep working instead of silently 400ing.
|
||
"""
|
||
import litellm
|
||
|
||
monkeypatch.setattr(
|
||
litellm,
|
||
"s3_callback_params",
|
||
{
|
||
"s3_bucket_name": "from-global",
|
||
"s3_server_side_encryption": "AES256",
|
||
"s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id",
|
||
},
|
||
)
|
||
logger = S3Logger()
|
||
assert logger.s3_server_side_encryption == "AES256"
|
||
assert logger.s3_sse_kms_key_id is None
|
||
|
||
|
||
def test_non_string_algorithm_is_dropped_and_valid_key_id_is_rescued(monkeypatch):
|
||
"""
|
||
A YAML boolean in s3_server_side_encryption must not crash logger init and
|
||
must not discard the valid key id; aws:kms is inferred from the key id.
|
||
"""
|
||
import litellm
|
||
|
||
monkeypatch.setattr(
|
||
litellm,
|
||
"s3_callback_params",
|
||
{
|
||
"s3_bucket_name": "from-global",
|
||
"s3_server_side_encryption": True,
|
||
"s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id",
|
||
},
|
||
)
|
||
logger = S3Logger()
|
||
assert logger.s3_server_side_encryption == "aws:kms"
|
||
assert logger.s3_sse_kms_key_id == ("arn:aws:kms:us-east-1:111122223333:key/test-key-id")
|
||
|
||
|
||
def test_non_string_key_id_is_dropped_and_valid_algorithm_is_kept(monkeypatch):
|
||
"""A mistyped key id (unquoted YAML number) must not disable the valid algorithm."""
|
||
import litellm
|
||
|
||
monkeypatch.setattr(
|
||
litellm,
|
||
"s3_callback_params",
|
||
{
|
||
"s3_bucket_name": "from-global",
|
||
"s3_server_side_encryption": "aws:kms",
|
||
"s3_sse_kms_key_id": 12345,
|
||
},
|
||
)
|
||
logger = S3Logger()
|
||
assert logger.s3_server_side_encryption == "aws:kms"
|
||
assert logger.s3_sse_kms_key_id is None
|
||
|
||
|
||
_ACCESS_KEY = "AKIAIOSFODNN7EXAMPLE"
|
||
_SECRET_KEY = "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"
|
||
_KEY_WITH_SPACE = "LOGS/LLM AI Projects/2026-08-04/time-13-01-00-abc.json"
|
||
|
||
|
||
def _signature_for(signer_cls, url: str, method: str, body: bytes | None, headers: dict[str, str]) -> str:
|
||
from botocore.awsrequest import AWSRequest
|
||
from botocore.credentials import Credentials
|
||
|
||
sent = {name.lower(): value for name, value in headers.items()}
|
||
signed_header_names = sent["authorization"].split("SignedHeaders=")[1].split(", ")[0].split(";")
|
||
request = AWSRequest(
|
||
method=method,
|
||
url=url,
|
||
data=body,
|
||
headers={name: sent[name] for name in signed_header_names if name in sent},
|
||
)
|
||
request.context["timestamp"] = sent["x-amz-date"]
|
||
signer = signer_cls(Credentials(_ACCESS_KEY, _SECRET_KEY), "s3", "us-east-1")
|
||
canonical_request = signer.canonical_request(request)
|
||
return signer.signature(signer.string_to_sign(request, canonical_request), request)
|
||
|
||
|
||
def _as_s3_canonicalizes(url: str) -> str:
|
||
"""
|
||
The path S3 rebuilds from the wire path: percent-encode everything outside the unreserved
|
||
set, without normalizing or double-encoding. `=` becomes `%3D`, `%20` stays `%20`.
|
||
"""
|
||
from urllib.parse import quote, unquote, urlsplit, urlunsplit
|
||
|
||
split = urlsplit(url)
|
||
return urlunsplit(split._replace(path=quote(unquote(split.path), safe="/~")))
|
||
|
||
|
||
def _assert_signed_for_s3_canonicalization(url: str, method: str, body: bytes | None, headers: dict[str, str]) -> None:
|
||
"""
|
||
S3 rebuilds the canonical request from the wire path with single percent-encoding, which
|
||
botocore models as S3SigV4Auth; plain SigV4Auth double-encodes it (%2520 for a space) and S3
|
||
answers 403 SignatureDoesNotMatch. Assert we sent an already-encoded path and signed it the
|
||
way S3 reads it.
|
||
"""
|
||
from botocore.auth import S3SigV4Auth, SigV4Auth
|
||
|
||
assert url == _as_s3_canonicalizes(url)
|
||
sent_signature = headers["Authorization"].split("Signature=")[1].strip()
|
||
assert sent_signature == _signature_for(S3SigV4Auth, url, method, body, headers)
|
||
assert sent_signature != _signature_for(SigV4Auth, url, method, body, headers)
|
||
|
||
|
||
def _logger_for_signing() -> S3Logger:
|
||
return S3Logger(
|
||
s3_bucket_name="logs-bucket",
|
||
s3_aws_access_key_id=_ACCESS_KEY,
|
||
s3_aws_secret_access_key=_SECRET_KEY,
|
||
s3_region_name="us-east-1",
|
||
)
|
||
|
||
|
||
def _element_with_space():
|
||
from litellm.types.integrations.s3_v2 import s3BatchLoggingElement
|
||
|
||
return s3BatchLoggingElement(
|
||
s3_object_key=_KEY_WITH_SPACE,
|
||
payload={"test": "sigv4"},
|
||
s3_object_download_filename="log.json",
|
||
)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_upload_signs_object_key_with_space_the_way_s3_does():
|
||
from unittest.mock import AsyncMock, MagicMock
|
||
|
||
logger = _logger_for_signing()
|
||
response = MagicMock()
|
||
response.status_code = 200
|
||
response.raise_for_status = MagicMock()
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put.return_value = response
|
||
|
||
await logger.async_upload_data_to_s3(_element_with_space())
|
||
|
||
call = logger.async_httpx_client.put.call_args
|
||
_assert_signed_for_s3_canonicalization(
|
||
url=call[0][0],
|
||
method="PUT",
|
||
body=call.kwargs["data"].encode("utf-8"),
|
||
headers=call.kwargs["headers"],
|
||
)
|
||
|
||
|
||
def test_sync_upload_signs_object_key_with_space_the_way_s3_does():
|
||
from unittest.mock import MagicMock
|
||
|
||
logger = _logger_for_signing()
|
||
response = MagicMock()
|
||
response.status_code = 200
|
||
response.raise_for_status = MagicMock()
|
||
mock_sync_client = MagicMock()
|
||
mock_sync_client.put.return_value = response
|
||
|
||
with patch("litellm.integrations.s3_v2.get_httpx_client", return_value=mock_sync_client):
|
||
logger.upload_data_to_s3(_element_with_space())
|
||
|
||
call = mock_sync_client.put.call_args
|
||
_assert_signed_for_s3_canonicalization(
|
||
url=call[0][0],
|
||
method="PUT",
|
||
body=call.kwargs["data"].encode("utf-8"),
|
||
headers=call.kwargs["headers"],
|
||
)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_download_signs_object_key_with_space_the_way_s3_does():
|
||
from unittest.mock import AsyncMock, MagicMock
|
||
|
||
logger = _logger_for_signing()
|
||
response = MagicMock()
|
||
response.status_code = 200
|
||
response.json = MagicMock(return_value={"downloaded": "data"})
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.get.return_value = response
|
||
|
||
assert await logger._download_object_from_s3(_KEY_WITH_SPACE) == {"downloaded": "data"}
|
||
|
||
call = logger.async_httpx_client.get.call_args
|
||
_assert_signed_for_s3_canonicalization(
|
||
url=call[0][0],
|
||
method="GET",
|
||
body=None,
|
||
headers=call.kwargs["headers"],
|
||
)
|
||
|
||
|
||
_RESERVED_CHAR_KEYS = (
|
||
"2026-08-21/time-05-29-36_resp_bGl0ZWxsbTpjdXN0b20=.json",
|
||
"session=logs/2026-08-21/time-05-29-36_abc.json",
|
||
"a+b/2026-08-21/time-05-29-36_abc.json",
|
||
"a&b/2026-08-21/time-05-29-36_abc.json",
|
||
"a#b/2026-08-21/time-05-29-36_abc.json",
|
||
"a?b/2026-08-21/time-05-29-36_abc.json",
|
||
"a%b/2026-08-21/time-05-29-36_abc.json",
|
||
_KEY_WITH_SPACE,
|
||
)
|
||
|
||
|
||
def _element_for(s3_object_key: str):
|
||
from litellm.types.integrations.s3_v2 import s3BatchLoggingElement
|
||
|
||
return s3BatchLoggingElement(
|
||
s3_object_key=s3_object_key,
|
||
payload={"test": "sigv4"},
|
||
s3_object_download_filename="log.json",
|
||
)
|
||
|
||
|
||
def _expected_wire_url(s3_object_key: str) -> str:
|
||
"""The URL boto3 itself would put on the wire for this key."""
|
||
from urllib.parse import quote
|
||
|
||
return f"https://logs-bucket.s3.us-east-1.amazonaws.com/{quote(s3_object_key, safe='/')}"
|
||
|
||
|
||
@pytest.mark.parametrize("s3_object_key", _RESERVED_CHAR_KEYS)
|
||
@pytest.mark.asyncio
|
||
async def test_async_upload_percent_encodes_reserved_characters_in_object_key(s3_object_key):
|
||
from unittest.mock import AsyncMock, MagicMock
|
||
|
||
logger = _logger_for_signing()
|
||
response = MagicMock()
|
||
response.status_code = 200
|
||
response.raise_for_status = MagicMock()
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put.return_value = response
|
||
|
||
await logger.async_upload_data_to_s3(_element_for(s3_object_key))
|
||
|
||
call = logger.async_httpx_client.put.call_args
|
||
assert call[0][0] == _expected_wire_url(s3_object_key)
|
||
_assert_signed_for_s3_canonicalization(
|
||
url=call[0][0],
|
||
method="PUT",
|
||
body=call.kwargs["data"].encode("utf-8"),
|
||
headers=call.kwargs["headers"],
|
||
)
|
||
|
||
|
||
@pytest.mark.parametrize("s3_object_key", _RESERVED_CHAR_KEYS)
|
||
def test_sync_upload_percent_encodes_reserved_characters_in_object_key(s3_object_key):
|
||
from unittest.mock import MagicMock
|
||
|
||
logger = _logger_for_signing()
|
||
response = MagicMock()
|
||
response.status_code = 200
|
||
response.raise_for_status = MagicMock()
|
||
mock_sync_client = MagicMock()
|
||
mock_sync_client.put.return_value = response
|
||
|
||
with patch("litellm.integrations.s3_v2.get_httpx_client", return_value=mock_sync_client):
|
||
logger.upload_data_to_s3(_element_for(s3_object_key))
|
||
|
||
call = mock_sync_client.put.call_args
|
||
assert call[0][0] == _expected_wire_url(s3_object_key)
|
||
_assert_signed_for_s3_canonicalization(
|
||
url=call[0][0],
|
||
method="PUT",
|
||
body=call.kwargs["data"].encode("utf-8"),
|
||
headers=call.kwargs["headers"],
|
||
)
|
||
|
||
|
||
@pytest.mark.parametrize("s3_object_key", _RESERVED_CHAR_KEYS)
|
||
@pytest.mark.asyncio
|
||
async def test_download_percent_encodes_reserved_characters_in_object_key(s3_object_key):
|
||
from unittest.mock import AsyncMock, MagicMock
|
||
|
||
logger = _logger_for_signing()
|
||
response = MagicMock()
|
||
response.status_code = 200
|
||
response.json = MagicMock(return_value={"downloaded": "data"})
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.get.return_value = response
|
||
|
||
assert await logger._download_object_from_s3(s3_object_key) == {"downloaded": "data"}
|
||
|
||
call = logger.async_httpx_client.get.call_args
|
||
assert call[0][0] == _expected_wire_url(s3_object_key)
|
||
_assert_signed_for_s3_canonicalization(
|
||
url=call[0][0],
|
||
method="GET",
|
||
body=None,
|
||
headers=call.kwargs["headers"],
|
||
)
|
||
|
||
|
||
def _s3_logger_for_region(region_name: str) -> S3Logger:
|
||
logger = S3Logger.__new__(S3Logger)
|
||
logger.s3_endpoint_url = None
|
||
logger.s3_bucket_name = "my-litellm-audit"
|
||
logger.s3_region_name = region_name
|
||
return logger
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
"region_name,expected_url",
|
||
[
|
||
(
|
||
"cn-northwest-1",
|
||
"https://my-litellm-audit.s3.cn-northwest-1.amazonaws.com.cn/2025-01-01/key.json",
|
||
),
|
||
(
|
||
"us-gov-west-1",
|
||
"https://my-litellm-audit.s3.us-gov-west-1.amazonaws.com/2025-01-01/key.json",
|
||
),
|
||
(
|
||
"us-east-1",
|
||
"https://my-litellm-audit.s3.us-east-1.amazonaws.com/2025-01-01/key.json",
|
||
),
|
||
],
|
||
)
|
||
def test_build_object_url_uses_partition_dns_suffix(region_name: str, expected_url: str) -> None:
|
||
assert _s3_logger_for_region(region_name)._build_object_url("2025-01-01/key.json") == expected_url
|
||
|
||
|
||
def _prompts_only_logger(s3_log_prompts_only: bool | None = None) -> S3Logger:
|
||
return S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_log_prompts_only=s3_log_prompts_only,
|
||
)
|
||
|
||
|
||
def _chat_payload() -> StandardLoggingPayload:
|
||
return StandardLoggingPayload(
|
||
id="chatcmpl-prompts-only",
|
||
messages=[{"role": "user", "content": "Reply with exactly the word PINEAPPLE."}],
|
||
response={"choices": [{"message": {"role": "assistant", "content": "PINEAPPLE"}}]},
|
||
metadata={"user_api_key_team_alias": None},
|
||
)
|
||
|
||
|
||
async def _queued_body_via_async_upload(
|
||
logger: S3Logger, log_event: Callable[..., Awaitable[None]]
|
||
) -> dict[str, object]:
|
||
payload = _chat_payload()
|
||
original = copy.deepcopy(payload)
|
||
await log_event(
|
||
kwargs={"standard_logging_object": payload},
|
||
response_obj=None,
|
||
start_time=datetime(2026, 7, 30, 12, 0, 0),
|
||
end_time=datetime(2026, 7, 30, 12, 0, 1),
|
||
)
|
||
assert payload == original, "the caller's standard_logging_object must not be mutated"
|
||
(element,) = logger.log_queue
|
||
|
||
response = MagicMock()
|
||
response.status_code = 200
|
||
response.raise_for_status = MagicMock()
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put.return_value = response
|
||
await logger.async_upload_data_to_s3(element)
|
||
return json.loads(logger.async_httpx_client.put.call_args.kwargs["data"])
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.parametrize("event_name", ["async_log_success_event", "async_log_failure_event"])
|
||
async def test_prompts_only_drops_response_but_keeps_messages_in_uploaded_object(
|
||
monkeypatch: pytest.MonkeyPatch, event_name: str
|
||
):
|
||
import litellm
|
||
|
||
monkeypatch.setattr(litellm, "s3_callback_params", {"s3_log_prompts_only": True})
|
||
logger = _prompts_only_logger()
|
||
|
||
log_event: Callable[..., Awaitable[None]] = (
|
||
logger.async_log_success_event if event_name == "async_log_success_event" else logger.async_log_failure_event
|
||
)
|
||
body = await _queued_body_via_async_upload(logger, log_event)
|
||
|
||
assert body["messages"] == _chat_payload()["messages"]
|
||
assert body["response"] is None
|
||
assert body["id"] == "chatcmpl-prompts-only"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_prompts_only_default_off_keeps_response_in_uploaded_object(monkeypatch: pytest.MonkeyPatch):
|
||
import litellm
|
||
|
||
monkeypatch.setattr(litellm, "s3_callback_params", {})
|
||
monkeypatch.delenv("S3_LOG_PROMPTS_ONLY", raising=False)
|
||
logger = _prompts_only_logger()
|
||
|
||
body = await _queued_body_via_async_upload(logger, logger.async_log_success_event)
|
||
|
||
assert body["response"] == _chat_payload()["response"]
|
||
assert body["messages"] == _chat_payload()["messages"]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_prompts_only_explicit_false_in_params_beats_env_var(monkeypatch: pytest.MonkeyPatch):
|
||
import litellm
|
||
|
||
monkeypatch.setattr(litellm, "s3_callback_params", {"s3_log_prompts_only": False})
|
||
monkeypatch.setenv("S3_LOG_PROMPTS_ONLY", "true")
|
||
logger = _prompts_only_logger()
|
||
|
||
body = await _queued_body_via_async_upload(logger, logger.async_log_success_event)
|
||
|
||
assert body["response"] == _chat_payload()["response"]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_prompts_only_env_var_applies_when_param_unset(monkeypatch: pytest.MonkeyPatch):
|
||
import litellm
|
||
|
||
monkeypatch.setattr(litellm, "s3_callback_params", {})
|
||
logger = _prompts_only_logger()
|
||
monkeypatch.setenv("S3_LOG_PROMPTS_ONLY", "true")
|
||
|
||
body = await _queued_body_via_async_upload(logger, logger.async_log_success_event)
|
||
|
||
assert body["response"] is None
|
||
assert body["messages"] == _chat_payload()["messages"]
|
||
|
||
|
||
@respx.mock
|
||
def test_prompts_only_constructor_kwarg_applies_to_sync_upload(monkeypatch: pytest.MonkeyPatch):
|
||
import litellm
|
||
|
||
monkeypatch.setattr(litellm, "s3_callback_params", {})
|
||
monkeypatch.delenv("S3_LOG_PROMPTS_ONLY", raising=False)
|
||
logger = _prompts_only_logger(s3_log_prompts_only=True)
|
||
payload = _chat_payload()
|
||
|
||
element = logger.create_s3_batch_logging_element(
|
||
start_time=datetime(2026, 7, 30, 12, 0, 0),
|
||
standard_logging_payload=payload,
|
||
)
|
||
assert element is not None
|
||
assert payload["response"] == _chat_payload()["response"]
|
||
|
||
put_route = respx.put(url__regex=r"https://test-bucket\.s3\..*").mock(return_value=httpx.Response(200))
|
||
logger.upload_data_to_s3(element)
|
||
|
||
body = json.loads(put_route.calls.last.request.content)
|
||
assert body["response"] is None
|
||
assert body["messages"] == _chat_payload()["messages"]
|
||
|
||
|
||
@pytest.mark.parametrize("callback_name", ["s3", "s3_v2"])
|
||
def test_prompts_only_toggle_is_exposed_to_admin_ui_for_both_s3_callbacks(callback_name: str):
|
||
from litellm.integrations.custom_logger import CustomLogger
|
||
|
||
assert "S3_LOG_PROMPTS_ONLY" in CustomLogger.get_callback_env_vars(callback_name)
|
||
|
||
|
||
_PARTITION_START: Final = datetime(2026, 9, 29, 14, 5, 9, 123456)
|
||
_PARTITION_ID: Final = "chatcmpl-partition"
|
||
|
||
|
||
def _partition_payload(response_id: str = _PARTITION_ID) -> StandardLoggingPayload:
|
||
return StandardLoggingPayload(
|
||
id=response_id,
|
||
metadata={"user_api_key_team_alias": "team-a", "user_api_key_alias": "key-a"},
|
||
messages=[],
|
||
)
|
||
|
||
|
||
def _partition_logger(
|
||
monkeypatch: pytest.MonkeyPatch, callback_params: dict[str, object], **kwargs: object
|
||
) -> S3Logger:
|
||
import litellm
|
||
|
||
monkeypatch.setattr(
|
||
litellm,
|
||
"s3_callback_params",
|
||
{"s3_bucket_name": "test-bucket", "s3_region_name": "us-east-1", "s3_path": "logs", **callback_params},
|
||
)
|
||
return S3Logger(
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_use_team_prefix=True,
|
||
s3_use_key_prefix=True,
|
||
**kwargs,
|
||
)
|
||
|
||
|
||
_DAILY_KEY: Final = f"logs/team-a/key-a/2026-09-29/time-14-05-09-123456_{_PARTITION_ID}.json"
|
||
_HOURLY_KEY: Final = f"logs/team-a/key-a/2026-09-29/14/time-14-05-09-123456_{_PARTITION_ID}.json"
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
("callback_params", "expected_key"),
|
||
[
|
||
({}, _DAILY_KEY),
|
||
({"s3_partition_granularity": None}, _DAILY_KEY),
|
||
({"s3_partition_granularity": "day"}, _DAILY_KEY),
|
||
({"s3_partition_granularity": "hour"}, _HOURLY_KEY),
|
||
],
|
||
)
|
||
def test_partition_granularity_sets_request_log_folder(
|
||
monkeypatch: pytest.MonkeyPatch, callback_params: dict[str, object], expected_key: str
|
||
) -> None:
|
||
monkeypatch.delenv("S3_PARTITION_GRANULARITY", raising=False)
|
||
logger = _partition_logger(monkeypatch, callback_params)
|
||
|
||
element = logger.create_s3_batch_logging_element(_PARTITION_START, _partition_payload())
|
||
|
||
assert element is not None
|
||
assert element.s3_object_key == expected_key
|
||
|
||
|
||
@pytest.mark.parametrize("invalid", ["hourly", "HOUR", "1", 1, True])
|
||
def test_invalid_partition_granularity_warns_and_keeps_daily_folder(
|
||
monkeypatch: pytest.MonkeyPatch, invalid: object
|
||
) -> None:
|
||
monkeypatch.delenv("S3_PARTITION_GRANULARITY", raising=False)
|
||
with patch("litellm.integrations.s3.verbose_logger") as mock_logger:
|
||
logger = _partition_logger(monkeypatch, {"s3_partition_granularity": invalid})
|
||
element = logger.create_s3_batch_logging_element(_PARTITION_START, _partition_payload())
|
||
second = logger.create_s3_batch_logging_element(_PARTITION_START, _partition_payload())
|
||
|
||
assert element is not None
|
||
assert second is not None
|
||
assert element.s3_object_key == second.s3_object_key == _DAILY_KEY
|
||
mock_logger.warning.assert_called_once()
|
||
assert mock_logger.warning.call_args.args[1:] == (invalid,)
|
||
|
||
|
||
def test_partition_granularity_reads_admin_ui_env_var_below_callback_params(monkeypatch: pytest.MonkeyPatch) -> None:
|
||
monkeypatch.setenv("S3_PARTITION_GRANULARITY", "hour")
|
||
|
||
from_env = _partition_logger(monkeypatch, {}).create_s3_batch_logging_element(
|
||
_PARTITION_START, _partition_payload()
|
||
)
|
||
from_params = _partition_logger(monkeypatch, {"s3_partition_granularity": "day"}).create_s3_batch_logging_element(
|
||
_PARTITION_START, _partition_payload()
|
||
)
|
||
|
||
assert from_env is not None and from_env.s3_object_key == _HOURLY_KEY
|
||
assert from_params is not None and from_params.s3_object_key == _DAILY_KEY
|
||
|
||
|
||
def test_partition_granularity_constructor_argument_and_os_environ_reference(monkeypatch: pytest.MonkeyPatch) -> None:
|
||
monkeypatch.delenv("S3_PARTITION_GRANULARITY", raising=False)
|
||
monkeypatch.setenv("MY_S3_PARTITION", "hour")
|
||
|
||
from_ctor = _partition_logger(monkeypatch, {}, s3_partition_granularity="hour")
|
||
from_secret = _partition_logger(monkeypatch, {"s3_partition_granularity": "os.environ/MY_S3_PARTITION"})
|
||
|
||
for logger in (from_ctor, from_secret):
|
||
element = logger.create_s3_batch_logging_element(_PARTITION_START, _partition_payload())
|
||
assert element is not None and element.s3_object_key == _HOURLY_KEY
|
||
|
||
|
||
def test_hourly_partition_long_key_keeps_hour_folder_within_s3_limit(monkeypatch: pytest.MonkeyPatch) -> None:
|
||
from litellm.constants import MAX_S3_OBJECT_KEY_BYTES
|
||
|
||
monkeypatch.delenv("S3_PARTITION_GRANULARITY", raising=False)
|
||
logger = _partition_logger(monkeypatch, {"s3_partition_granularity": "hour", "s3_path": "p" * 1100})
|
||
|
||
element = logger.create_s3_batch_logging_element(_PARTITION_START, _partition_payload("r" * 600))
|
||
|
||
assert element is not None
|
||
assert len(element.s3_object_key.encode("utf-8")) <= MAX_S3_OBJECT_KEY_BYTES
|
||
assert re.search(r"/2026-09-29/14/[0-9a-f]{64}\.json$", element.s3_object_key)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.parametrize(("granularity", "hour_folder"), [("hour", True), ("day", False), (None, False)])
|
||
async def test_audit_log_key_follows_audit_callback_params_partition_granularity(
|
||
monkeypatch: pytest.MonkeyPatch, granularity: str | None, hour_folder: bool
|
||
) -> None:
|
||
monkeypatch.delenv("S3_PARTITION_GRANULARITY", raising=False)
|
||
logger = S3Logger(
|
||
s3_callback_params_override={
|
||
"s3_bucket_name": "audit-bucket",
|
||
"s3_path": "audit",
|
||
"s3_partition_granularity": granularity,
|
||
}
|
||
)
|
||
|
||
await logger.async_log_audit_log_event({"id": "audit-1"})
|
||
|
||
(element,) = logger.log_queue
|
||
match = re.fullmatch(
|
||
r"audit/audit_logs/\d{4}-\d{2}-\d{2}/(?:(\d{2})/)?(\d{2})-\d{2}-\d{2}_audit-1\.json", element.s3_object_key
|
||
)
|
||
assert match is not None, element.s3_object_key
|
||
assert (match.group(1) is not None) is hour_folder
|
||
if hour_folder:
|
||
assert match.group(1) == match.group(2)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_hourly_batch_file_upload_writes_one_file_per_hour_folder(monkeypatch: pytest.MonkeyPatch) -> None:
|
||
monkeypatch.delenv("S3_PARTITION_GRANULARITY", raising=False)
|
||
logger = _partition_logger(monkeypatch, {"s3_partition_granularity": "hour"}, s3_batch_file_upload=True)
|
||
put = _RecordingPut()
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
before = logger.create_s3_batch_logging_element(datetime(2026, 9, 29, 13, 59, 59), _partition_payload("before"))
|
||
after = logger.create_s3_batch_logging_element(datetime(2026, 9, 29, 14, 0, 1), _partition_payload("after"))
|
||
assert before is not None and after is not None
|
||
logger.log_queue = [before, after]
|
||
|
||
await logger.async_send_batch()
|
||
|
||
by_folder = {
|
||
re.sub(r"/batch_\d{2}-\d{2}-\d{2}_[0-9a-f]{32}\.jsonl$", "", url.split(".com/", 1)[-1]): data
|
||
for url, data, _headers in put.calls
|
||
}
|
||
assert sorted(by_folder) == ["logs/team-a/key-a/2026-09-29/13", "logs/team-a/key-a/2026-09-29/14"]
|
||
assert [json.loads(line)["id"] for line in (by_folder["logs/team-a/key-a/2026-09-29/13"] or "").splitlines()] == [
|
||
"before"
|
||
]
|
||
assert [json.loads(line)["id"] for line in (by_folder["logs/team-a/key-a/2026-09-29/14"] or "").splitlines()] == [
|
||
"after"
|
||
]
|
||
|
||
|
||
@pytest.mark.parametrize("granularity", [None, "day", "hour"])
|
||
def test_cold_storage_object_key_matches_the_uploaded_request_log_key(
|
||
monkeypatch: pytest.MonkeyPatch, granularity: str | None
|
||
) -> None:
|
||
import litellm
|
||
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
|
||
|
||
monkeypatch.delenv("S3_PARTITION_GRANULARITY", raising=False)
|
||
monkeypatch.setattr(
|
||
litellm,
|
||
"s3_callback_params",
|
||
{"s3_bucket_name": "test-bucket", "s3_path": "coldlogs", "s3_partition_granularity": granularity},
|
||
)
|
||
monkeypatch.setattr(litellm, "cold_storage_custom_logger", "s3_v2")
|
||
logger = S3Logger()
|
||
uploaded = logger.create_s3_batch_logging_element(
|
||
_PARTITION_START, StandardLoggingPayload(id=_PARTITION_ID, metadata={}, messages=[])
|
||
)
|
||
|
||
monkeypatch.setattr(litellm, "callbacks", [logger])
|
||
cold_key = StandardLoggingPayloadSetup._generate_cold_storage_object_key(
|
||
start_time=_PARTITION_START, response_id=_PARTITION_ID
|
||
)
|
||
|
||
assert uploaded is not None
|
||
assert cold_key == uploaded.s3_object_key
|
||
assert ("/2026-09-29/14/" in cold_key) is (granularity == "hour")
|
||
|
||
|
||
def test_cold_storage_key_matches_upload_when_env_var_changes_mid_request(monkeypatch: pytest.MonkeyPatch) -> None:
|
||
import litellm
|
||
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
|
||
|
||
monkeypatch.delenv("S3_PARTITION_GRANULARITY", raising=False)
|
||
monkeypatch.setattr(litellm, "s3_callback_params", {"s3_bucket_name": "test-bucket", "s3_path": "coldlogs"})
|
||
monkeypatch.setattr(litellm, "cold_storage_custom_logger", "s3_v2")
|
||
logger = S3Logger()
|
||
monkeypatch.setattr(litellm, "callbacks", [logger])
|
||
|
||
cold_key = StandardLoggingPayloadSetup._generate_cold_storage_object_key(
|
||
start_time=_PARTITION_START, response_id=_PARTITION_ID
|
||
)
|
||
monkeypatch.setenv("S3_PARTITION_GRANULARITY", "hour")
|
||
uploaded = logger.create_s3_batch_logging_element(
|
||
_PARTITION_START,
|
||
StandardLoggingPayload(id=_PARTITION_ID, metadata={"cold_storage_object_key": cold_key}, messages=[]),
|
||
)
|
||
|
||
assert uploaded is not None
|
||
assert cold_key == uploaded.s3_object_key == f"coldlogs/2026-09-29/time-14-05-09-123456_{_PARTITION_ID}.json"
|
||
|
||
|
||
def test_hour_upload_ignores_a_cold_storage_key_owned_by_another_logger(monkeypatch: pytest.MonkeyPatch) -> None:
|
||
import litellm
|
||
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
|
||
|
||
monkeypatch.delenv("S3_PARTITION_GRANULARITY", raising=False)
|
||
monkeypatch.setattr(
|
||
litellm, "s3_callback_params", {"s3_bucket_name": "test-bucket", "s3_partition_granularity": "hour"}
|
||
)
|
||
monkeypatch.setattr(litellm, "cold_storage_custom_logger", "gcs_bucket")
|
||
logger = S3Logger()
|
||
cold_key = StandardLoggingPayloadSetup._generate_cold_storage_object_key(
|
||
start_time=_PARTITION_START, response_id=_PARTITION_ID
|
||
)
|
||
uploaded = logger.create_s3_batch_logging_element(
|
||
_PARTITION_START,
|
||
StandardLoggingPayload(id=_PARTITION_ID, metadata={"cold_storage_object_key": cold_key}, messages=[]),
|
||
)
|
||
|
||
assert cold_key == f"2026-09-29/time-14-05-09-123456_{_PARTITION_ID}.json"
|
||
assert uploaded is not None
|
||
assert uploaded.s3_object_key == f"2026-09-29/14/time-14-05-09-123456_{_PARTITION_ID}.json"
|
||
|
||
|
||
@pytest.mark.parametrize("callback_name", ["s3", "s3_v2"])
|
||
def test_partition_granularity_is_exposed_to_admin_ui(callback_name: str) -> None:
|
||
from litellm.integrations.custom_logger import CustomLogger
|
||
|
||
assert "S3_PARTITION_GRANULARITY" in CustomLogger.get_callback_env_vars(callback_name)
|
||
|
||
|
||
def _element(payload: dict[str, object], key_suffix: str) -> s3BatchLoggingElement:
|
||
return s3BatchLoggingElement(
|
||
s3_object_key=f"2025-09-14/test-{key_suffix}.json",
|
||
payload=payload,
|
||
s3_object_download_filename=f"test-{key_suffix}.json",
|
||
)
|
||
|
||
|
||
def _ok_response() -> MagicMock:
|
||
response = MagicMock()
|
||
response.status_code = 200
|
||
response.text = ""
|
||
response.raise_for_status = MagicMock()
|
||
return response
|
||
|
||
|
||
class _CountingPut:
|
||
def __init__(self, width: int) -> None:
|
||
self.width = width
|
||
self.in_flight = 0
|
||
self.peak = 0
|
||
self.calls = 0
|
||
|
||
async def __call__(
|
||
self,
|
||
url: str,
|
||
data: str | None = None,
|
||
headers: dict[str, str] | None = None,
|
||
timeout: httpx.Timeout | None = None,
|
||
) -> MagicMock:
|
||
self.in_flight += 1
|
||
self.peak = max(self.peak, self.in_flight)
|
||
self.calls += 1
|
||
for _ in range(50):
|
||
if self.in_flight >= self.width:
|
||
break
|
||
await _real_sleep(0)
|
||
self.in_flight -= 1
|
||
return _ok_response()
|
||
|
||
|
||
class _RecordingPut:
|
||
def __init__(self) -> None:
|
||
self.calls: tuple[tuple[str, str | None, dict[str, str] | None], ...] = ()
|
||
|
||
async def __call__(
|
||
self,
|
||
url: str,
|
||
data: str | None = None,
|
||
headers: dict[str, str] | None = None,
|
||
timeout: httpx.Timeout | None = None,
|
||
) -> MagicMock:
|
||
self.calls = (*self.calls, (url, data, headers))
|
||
return _ok_response()
|
||
|
||
|
||
class _LateAppendingPut:
|
||
def __init__(self, logger: S3Logger, element: s3BatchLoggingElement, fail_first: bool = False) -> None:
|
||
self.logger = logger
|
||
self.element = element
|
||
self.fail_first = fail_first
|
||
self.appended = False
|
||
self.failed_key: str | None = None
|
||
|
||
async def __call__(
|
||
self,
|
||
url: str,
|
||
data: str | None = None,
|
||
headers: dict[str, str] | None = None,
|
||
timeout: httpx.Timeout | None = None,
|
||
) -> MagicMock:
|
||
if not self.appended:
|
||
self.appended = True
|
||
self.logger.log_queue.append(self.element)
|
||
if self.fail_first:
|
||
self.failed_key = url
|
||
if url == self.failed_key:
|
||
return _transient_failure_response()
|
||
return _ok_response()
|
||
|
||
|
||
class _AppendingFailingPut:
|
||
def __init__(self, logger: S3Logger, elements: tuple[s3BatchLoggingElement, ...]) -> None:
|
||
self.logger = logger
|
||
self.elements = elements
|
||
self.appended = False
|
||
|
||
async def __call__(
|
||
self,
|
||
url: str,
|
||
data: str | None = None,
|
||
headers: dict[str, str] | None = None,
|
||
timeout: httpx.Timeout | None = None,
|
||
) -> MagicMock:
|
||
if not self.appended:
|
||
self.appended = True
|
||
for element in self.elements:
|
||
self.logger.log_queue.append(element)
|
||
return _transient_failure_response()
|
||
|
||
|
||
class _FailOnSuffixPut:
|
||
def __init__(self, suffixes: tuple[str, ...]) -> None:
|
||
self.failing = True
|
||
self.suffixes = suffixes
|
||
self.calls: tuple[str, ...] = ()
|
||
|
||
async def __call__(
|
||
self,
|
||
url: str,
|
||
data: str | None = None,
|
||
headers: dict[str, str] | None = None,
|
||
timeout: httpx.Timeout | None = None,
|
||
) -> MagicMock:
|
||
self.calls = (*self.calls, url)
|
||
if self.failing and url.endswith(self.suffixes):
|
||
return _transient_failure_response()
|
||
return _ok_response()
|
||
|
||
|
||
class _FailUntilClearedPut:
|
||
def __init__(self, status: int = 503, code: str | None = "SlowDown", raw_body: str | None = None) -> None:
|
||
self.failing = True
|
||
self.response: Final = _coded_failure_response(status, code, raw_body)
|
||
self.calls: tuple[tuple[str, str | None], ...] = ()
|
||
|
||
async def __call__(
|
||
self,
|
||
url: str,
|
||
data: str | None = None,
|
||
headers: dict[str, str] | None = None,
|
||
timeout: httpx.Timeout | None = None,
|
||
) -> MagicMock:
|
||
self.calls = (*self.calls, (url, data))
|
||
if self.failing:
|
||
return self.response
|
||
return _ok_response()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_send_batch_bounds_concurrent_uploads() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_max_concurrent_uploads=4,
|
||
)
|
||
|
||
put = _CountingPut(logger.s3_max_concurrent_uploads)
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
|
||
logger.log_queue = [_element({"i": i}, f"{i}") for i in range(40)]
|
||
|
||
await logger.async_send_batch()
|
||
|
||
assert put.peak == 4
|
||
assert put.calls == 40
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_send_batch_uploads_single_jsonl_file() -> None:
|
||
import json
|
||
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_batch_file_upload=True,
|
||
)
|
||
|
||
put = _RecordingPut()
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
|
||
payloads = [{"id": "req-1"}, {"id": "req-2"}, {"id": "req-3"}]
|
||
logger.log_queue = [_element(payload, f"{i}") for i, payload in enumerate(payloads)]
|
||
|
||
await logger.async_send_batch()
|
||
|
||
assert len(put.calls) == 1
|
||
url, data, headers = put.calls[0]
|
||
assert url.endswith(".jsonl")
|
||
assert data is not None
|
||
assert headers is not None
|
||
assert [json.loads(line) for line in data.splitlines()] == payloads
|
||
assert headers["Content-Type"] == "application/x-ndjson"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_flush_queue_preserves_events_added_during_upload() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
|
||
late_element = _element({"id": "late"}, "late")
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = _LateAppendingPut(logger, late_element)
|
||
|
||
logger.log_queue = [_element({"id": "first"}, "first")]
|
||
|
||
await logger.flush_queue()
|
||
|
||
assert logger.log_queue == [late_element]
|
||
|
||
|
||
def _override_logger(**overrides: object) -> S3Logger:
|
||
return S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_callback_params_override=overrides,
|
||
)
|
||
|
||
|
||
def test_env_backed_false_string_keeps_per_request_uploads() -> None:
|
||
assert _override_logger(s3_batch_file_upload="false").s3_batch_file_upload is False
|
||
assert _override_logger(s3_batch_file_upload="true").s3_batch_file_upload is True
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_batch_file_upload=True,
|
||
s3_callback_params_override={"s3_batch_file_upload": "false"},
|
||
)
|
||
assert logger.s3_batch_file_upload is True
|
||
|
||
|
||
@pytest.mark.parametrize("bad", [0, -3, "0", "abc", ""])
|
||
def test_invalid_concurrency_falls_back_to_default(bad: object) -> None:
|
||
from litellm.constants import DEFAULT_S3_MAX_CONCURRENT_UPLOADS
|
||
|
||
logger = _override_logger(s3_max_concurrent_uploads=bad)
|
||
|
||
assert logger.s3_max_concurrent_uploads == DEFAULT_S3_MAX_CONCURRENT_UPLOADS
|
||
assert logger._upload_limiter._value == DEFAULT_S3_MAX_CONCURRENT_UPLOADS
|
||
|
||
|
||
def test_env_backed_concurrency_string_is_parsed() -> None:
|
||
logger = _override_logger(s3_max_concurrent_uploads="4")
|
||
|
||
assert logger.s3_max_concurrent_uploads == 4
|
||
assert logger._upload_limiter._value == 4
|
||
|
||
|
||
@pytest.mark.parametrize("empty", [None, ""])
|
||
def test_empty_config_concurrency_falls_back_to_constructor_value(empty: object) -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_max_concurrent_uploads=4,
|
||
s3_callback_params_override={"s3_max_concurrent_uploads": empty},
|
||
)
|
||
|
||
assert logger.s3_max_concurrent_uploads == 4
|
||
assert logger._upload_limiter._value == 4
|
||
|
||
|
||
def _coded_failure_response(status: int, code: str | None, raw_body: str | None = None) -> MagicMock:
|
||
body: Final = (
|
||
raw_body if raw_body is not None else (f"<Error><Code>{code}</Code></Error>" if code is not None else "")
|
||
)
|
||
response = MagicMock()
|
||
response.status_code = status
|
||
response.text = body
|
||
response.raise_for_status = MagicMock(
|
||
side_effect=httpx.HTTPStatusError(str(status), request=MagicMock(), response=response)
|
||
)
|
||
return response
|
||
|
||
|
||
def _transient_failure_response(status: int = 503) -> MagicMock:
|
||
return _coded_failure_response(status, "SlowDown")
|
||
|
||
|
||
def _terminal_failure_response() -> MagicMock:
|
||
return _coded_failure_response(400, "EntityTooLarge")
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_failed_uploads_stay_queued_for_next_flush() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
|
||
elements = [_element({"i": i}, f"{i}") for i in range(5)]
|
||
put = _FailOnSuffixPut(("test-2.json", "test-4.json"))
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
logger.log_queue = list(elements)
|
||
|
||
with patch("asyncio.sleep", new_callable=AsyncMock):
|
||
await logger.flush_queue()
|
||
|
||
assert [element.s3_object_key for element in logger.log_queue] == [
|
||
elements[2].s3_object_key,
|
||
elements[4].s3_object_key,
|
||
]
|
||
|
||
put.failing = False
|
||
await logger.flush_queue()
|
||
|
||
assert logger.log_queue == []
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_batch_file_upload_failure_keeps_whole_batch() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_batch_file_upload=True,
|
||
)
|
||
|
||
put = _FailUntilClearedPut()
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
|
||
elements = [_element({"i": i}, f"{i}") for i in range(3)]
|
||
logger.log_queue = list(elements)
|
||
|
||
with patch("asyncio.sleep", new_callable=AsyncMock):
|
||
await logger.flush_queue()
|
||
|
||
assert len(put.calls) == 3
|
||
assert len(logger.log_queue) == 1
|
||
assert logger.log_queue[0].body == "\n".join(json.dumps(element.payload) for element in elements)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_events_appended_during_failed_flush_survive() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
|
||
late = _element({"id": "late"}, "late")
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = _LateAppendingPut(logger, late, fail_first=True)
|
||
|
||
first = _element({"id": "first"}, "first")
|
||
logger.log_queue = [first]
|
||
|
||
with patch("asyncio.sleep", new_callable=AsyncMock):
|
||
await logger.flush_queue()
|
||
|
||
assert [element.s3_object_key for element in logger.log_queue] == [first.s3_object_key, late.s3_object_key]
|
||
assert logger.log_queue[0].retrying_since is None
|
||
|
||
logger.async_httpx_client.put.failed_key = None
|
||
await logger.flush_queue()
|
||
|
||
assert logger.log_queue == []
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_batch_file_key_shape() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_path="logs",
|
||
s3_batch_file_upload=True,
|
||
)
|
||
|
||
put = _RecordingPut()
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
logger.log_queue = [_element({"id": "req-1"}, "0")]
|
||
|
||
await logger.async_send_batch()
|
||
|
||
((url, _data, headers),) = put.calls
|
||
assert headers is not None
|
||
assert re.search(r".*/2025-09-14/batch_\d{2}-\d{2}-\d{2}_[0-9a-f]{32}\.jsonl$", url)
|
||
assert headers["Content-Disposition"].endswith('.jsonl"')
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_batch_file_groups_raw_elements_by_key_parent() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_batch_file_upload=True,
|
||
)
|
||
|
||
put = _RecordingPut()
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
|
||
alpha = s3BatchLoggingElement(
|
||
s3_object_key="logs/alpha/2026-01-01/a.json", payload={"id": "a"}, s3_object_download_filename="a.json"
|
||
)
|
||
beta = s3BatchLoggingElement(
|
||
s3_object_key="logs/beta/2026-01-01/b.json", payload={"id": "b"}, s3_object_download_filename="b.json"
|
||
)
|
||
plain = s3BatchLoggingElement(
|
||
s3_object_key="logs/2026-01-01/c.json", payload={"id": "c"}, s3_object_download_filename="c.json"
|
||
)
|
||
root = s3BatchLoggingElement(
|
||
s3_object_key="solo.json", payload={"id": "d"}, s3_object_download_filename="solo.json"
|
||
)
|
||
logger.log_queue = [alpha, beta, plain, root]
|
||
|
||
await logger.async_send_batch()
|
||
|
||
assert len(put.calls) == 4
|
||
by_parent = {
|
||
re.sub(r"(^|/)batch_\d{2}-\d{2}-\d{2}_[0-9a-f]{32}\.jsonl$", "", url.split(".com/", 1)[-1]): (url, data)
|
||
for url, data, _headers in put.calls
|
||
}
|
||
assert sorted(by_parent) == ["", "logs/2026-01-01", "logs/alpha/2026-01-01", "logs/beta/2026-01-01"]
|
||
assert [line for line in by_parent[""][1].splitlines()] == [json.dumps({"id": "d"})]
|
||
assert [line for line in by_parent["logs/alpha/2026-01-01"][1].splitlines()] == [json.dumps({"id": "a"})]
|
||
assert [line for line in by_parent["logs/beta/2026-01-01"][1].splitlines()] == [json.dumps({"id": "b"})]
|
||
assert [line for line in by_parent["logs/2026-01-01"][1].splitlines()] == [json.dumps({"id": "c"})]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_failed_batch_file_is_requeued_and_resent_unchanged() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_batch_file_upload=True,
|
||
)
|
||
|
||
put = _FailUntilClearedPut()
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
logger.log_queue = [_element({"i": i}, f"{i}") for i in range(3)]
|
||
|
||
with patch("asyncio.sleep", new_callable=AsyncMock):
|
||
await logger.flush_queue()
|
||
|
||
assert len(logger.log_queue) == 1
|
||
assert logger.log_queue[0].body is not None
|
||
assert logger.log_queue[0].s3_object_key.endswith(".jsonl")
|
||
|
||
put.failing = False
|
||
await logger.flush_queue()
|
||
|
||
assert logger.log_queue == []
|
||
assert len(put.calls) == 4
|
||
assert put.calls[0] == put.calls[1]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_elements_appended_after_failed_batch_file_get_their_own_file() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_batch_file_upload=True,
|
||
)
|
||
|
||
put = _FailUntilClearedPut()
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
logger.log_queue = [_element({"id": "first"}, "first")]
|
||
|
||
with patch("asyncio.sleep", new_callable=AsyncMock):
|
||
await logger.flush_queue()
|
||
|
||
late = _element({"id": "late"}, "late")
|
||
logger.log_queue.append(late)
|
||
|
||
put.failing = False
|
||
await logger.flush_queue()
|
||
|
||
assert logger.log_queue == []
|
||
assert len(put.calls) == 5
|
||
assert put.calls[0] == put.calls[1]
|
||
second_flush: Final = put.calls[3:]
|
||
assert put.calls[0] in second_flush
|
||
late_call: Final = next(call for call in second_flush if call != put.calls[0])
|
||
assert late_call[0] != put.calls[0][0]
|
||
assert late_call[1] == json.dumps({"id": "late"})
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_batch_file_mode_disabled_when_s3_v2_is_cold_storage_logger(monkeypatch: pytest.MonkeyPatch) -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_batch_file_upload=True,
|
||
)
|
||
|
||
put = _RecordingPut()
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
|
||
import litellm
|
||
|
||
monkeypatch.setattr(litellm, "cold_storage_custom_logger", "s3_v2")
|
||
logger.log_queue = [_element({"id": "req-1"}, "0")]
|
||
|
||
await logger.async_send_batch()
|
||
|
||
assert len(put.calls) == 1
|
||
assert put.calls[0][0].endswith("test-0.json")
|
||
|
||
monkeypatch.setattr(litellm, "cold_storage_custom_logger", None)
|
||
logger.log_queue = [_element({"id": "req-2"}, "1")]
|
||
|
||
await logger.async_send_batch()
|
||
|
||
assert len(put.calls) == 2
|
||
assert put.calls[1][0].endswith(".jsonl")
|
||
|
||
|
||
class _FailOnSuffixCodedPut:
|
||
def __init__(
|
||
self, suffixes: tuple[str, ...], status: int, code: str | None = None, raw_body: str | None = None
|
||
) -> None:
|
||
self.suffixes = suffixes
|
||
self.response: Final = _coded_failure_response(status, code, raw_body)
|
||
self.calls: tuple[str, ...] = ()
|
||
|
||
async def __call__(
|
||
self,
|
||
url: str,
|
||
data: str | None = None,
|
||
headers: dict[str, str] | None = None,
|
||
timeout: httpx.Timeout | None = None,
|
||
) -> MagicMock:
|
||
self.calls = (*self.calls, url)
|
||
if url.endswith(self.suffixes):
|
||
return self.response
|
||
return _ok_response()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.parametrize(
|
||
("status", "code", "raw_body", "puts_per_element"),
|
||
[
|
||
pytest.param(403, "AccessDenied", None, 3, id="access-denied-403"),
|
||
pytest.param(403, None, None, 3, id="empty-403"),
|
||
pytest.param(403, None, "<html>Forbidden</html>", 3, id="html-403"),
|
||
pytest.param(400, "KMS.DisabledException", None, 1, id="kms-disabled-400"),
|
||
pytest.param(404, "NoSuchBucket", None, 1, id="no-such-bucket-404"),
|
||
],
|
||
)
|
||
async def test_non_terminal_failure_is_requeued_and_delivered_on_recovery(
|
||
status: int, code: str | None, raw_body: str | None, puts_per_element: int
|
||
) -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
|
||
elements = [_element({"i": i}, f"{i}") for i in range(5)]
|
||
put = _FailUntilClearedPut(status=status, code=code, raw_body=raw_body)
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
logger.log_queue = list(elements)
|
||
|
||
with patch("asyncio.sleep", new_callable=AsyncMock):
|
||
await logger.flush_queue()
|
||
|
||
assert len(logger.log_queue) == 5
|
||
assert len(put.calls) == 5 * puts_per_element
|
||
|
||
put.failing = False
|
||
await logger.flush_queue()
|
||
|
||
assert logger.log_queue == []
|
||
assert len(put.calls) == 5 * puts_per_element + 5
|
||
landed: Final = frozenset(
|
||
element.s3_object_key
|
||
for element in elements
|
||
if any(call[0].endswith(element.s3_object_key) for call in put.calls[-5:])
|
||
)
|
||
assert landed == frozenset(element.s3_object_key for element in elements)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_persistent_500_stays_queued_through_a_dozen_failed_flushes() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
|
||
put = _FailUntilClearedPut(status=500, code="InternalError")
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
logger.log_queue = [_element({"i": i}, f"{i}") for i in range(5)]
|
||
|
||
with patch("asyncio.sleep", new_callable=AsyncMock):
|
||
for _ in range(12):
|
||
await logger.flush_queue()
|
||
assert len(logger.log_queue) == 5
|
||
|
||
assert len(put.calls) == 12 * 15
|
||
|
||
put.failing = False
|
||
await logger.flush_queue()
|
||
|
||
assert logger.log_queue == []
|
||
assert len(put.calls) == 12 * 15 + 5
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_terminal_object_is_dropped_once_next_to_delivered_siblings_when_opted_in() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_drop_on_terminal_error=True,
|
||
)
|
||
|
||
elements = [_element({"i": i}, f"{i}") for i in range(5)]
|
||
put = _FailOnSuffixCodedPut(("test-1.json",), 400, "EntityTooLarge")
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
logger.log_queue = list(elements)
|
||
|
||
await logger.flush_queue()
|
||
|
||
assert len(put.calls) == 5
|
||
assert logger.log_queue == []
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_terminal_object_is_requeued_when_opted_out() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_drop_on_terminal_error=False,
|
||
)
|
||
|
||
put = _FailOnSuffixCodedPut(("test-1.json",), 400, "EntityTooLarge")
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
logger.log_queue = [_element({"i": i}, f"{i}") for i in range(5)]
|
||
|
||
await logger.flush_queue()
|
||
|
||
assert len(logger.log_queue) == 1
|
||
assert logger.log_queue[0].s3_object_key.endswith("test-1.json")
|
||
assert sum(call.endswith("test-1.json") for call in put.calls) == 1
|
||
assert len(put.calls) == 5
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_terminal_objects_are_requeued_when_every_upload_in_the_flush_fails() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_drop_on_terminal_error=True,
|
||
)
|
||
|
||
put = _FailUntilClearedPut(status=400, code="EntityTooLarge")
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
logger.log_queue = [_element({"i": i}, f"{i}") for i in range(5)]
|
||
|
||
await logger.flush_queue()
|
||
|
||
assert len(logger.log_queue) == 5
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_retrying_past_the_opted_in_budget_is_dropped_only_next_to_delivered_siblings() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_max_retry_age_seconds=60,
|
||
)
|
||
|
||
aged = _element({"id": "aged"}, "aged").model_copy(update={"retrying_since": _NOW - 120})
|
||
fresh = _element({"id": "fresh"}, "fresh")
|
||
put = _FailOnSuffixCodedPut(("test-aged.json",), 503, "SlowDown")
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
logger.log_queue = [aged, fresh]
|
||
|
||
with (
|
||
patch("asyncio.sleep", new_callable=AsyncMock),
|
||
patch("time.monotonic", return_value=_NOW),
|
||
):
|
||
await logger.flush_queue()
|
||
|
||
assert logger.log_queue == []
|
||
assert len(put.calls) == 4
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_retrying_past_the_budget_stays_queued_when_the_whole_flush_fails() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_max_retry_age_seconds=60,
|
||
)
|
||
|
||
aged = _element({"id": "aged"}, "aged").model_copy(update={"retrying_since": _NOW - 120})
|
||
put = _FailUntilClearedPut(status=503, code="SlowDown")
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
logger.log_queue = [aged]
|
||
|
||
with (
|
||
patch("asyncio.sleep", new_callable=AsyncMock),
|
||
patch("time.monotonic", return_value=_NOW),
|
||
):
|
||
await logger.flush_queue()
|
||
|
||
assert len(logger.log_queue) == 1
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_overflow_after_a_failed_flush_trims_failed_first_and_counts_upload_failures_only(
|
||
caplog,
|
||
) -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_max_queue_size=4,
|
||
)
|
||
|
||
late = tuple(_element({"id": f"late-{index}"}, f"late-{index}") for index in range(3))
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = _AppendingFailingPut(logger, late)
|
||
logger.log_queue = [_element({"id": "first"}, "first"), _element({"id": "second"}, "second")]
|
||
|
||
with (
|
||
patch.object(logger, "handle_callback_failure") as mock_failure,
|
||
patch("asyncio.sleep", new_callable=AsyncMock),
|
||
pytest.raises(S3BatchUploadError),
|
||
):
|
||
await logger.async_send_batch()
|
||
|
||
assert [element.payload["id"] for element in logger.log_queue] == ["second", "late-0", "late-1", "late-2"]
|
||
failed_uploads: Final = 2
|
||
assert mock_failure.call_count == failed_uploads
|
||
mock_failure.assert_called_with(callback_name="S3Logger")
|
||
assert "dropped 1 oldest events" in caplog.text
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_default_logger_ages_out_elements_retrying_longer_than_an_hour(caplog) -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
|
||
elements = [
|
||
_element({"i": index}, f"{index}").model_copy(update={"retrying_since": _NOW - 7200}) for index in range(3)
|
||
]
|
||
put = _FailOnSuffixPut(("test-1.json", "test-2.json"))
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
logger.log_queue = list(elements)
|
||
|
||
with (
|
||
patch("asyncio.sleep", new_callable=AsyncMock),
|
||
patch("time.monotonic", return_value=_NOW),
|
||
):
|
||
await logger.flush_queue()
|
||
|
||
assert logger.log_queue == []
|
||
assert "uploads dropped" in caplog.text
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_opted_out_logger_never_ages_out_long_retrying_elements(caplog) -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_max_retry_age_seconds=0,
|
||
)
|
||
|
||
elements = [
|
||
_element({"i": index}, f"{index}").model_copy(update={"retrying_since": _NOW - 7200}) for index in range(3)
|
||
]
|
||
put = _FailOnSuffixPut(("test-1.json", "test-2.json"))
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
logger.log_queue = list(elements)
|
||
|
||
with (
|
||
patch("asyncio.sleep", new_callable=AsyncMock),
|
||
patch("time.monotonic", return_value=_NOW),
|
||
):
|
||
await logger.flush_queue()
|
||
|
||
assert [element.s3_object_key for element in logger.log_queue] == [
|
||
elements[1].s3_object_key,
|
||
elements[2].s3_object_key,
|
||
]
|
||
assert "uploads dropped" not in caplog.text
|
||
|
||
put.failing = False
|
||
await logger.flush_queue()
|
||
|
||
assert logger.log_queue == []
|
||
landed: Final = frozenset(call_url.rsplit("/", 1)[-1] for call_url in put.calls)
|
||
assert landed == frozenset(f"test-{index}.json" for index in range(3))
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_queue_grows_past_the_cap_while_the_sink_fails_and_everything_lands() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_max_queue_size=5,
|
||
)
|
||
|
||
elements = [_element({"i": index}, f"{index}") for index in range(8)]
|
||
put = _FailUntilClearedPut()
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
logger.log_queue = list(elements[:5])
|
||
|
||
with (
|
||
patch.object(logger, "handle_callback_failure") as mock_failure,
|
||
patch("asyncio.sleep", new_callable=AsyncMock),
|
||
):
|
||
await logger.flush_queue()
|
||
|
||
assert len(logger.log_queue) == 5
|
||
upload_failures: Final = 5
|
||
assert mock_failure.call_count == upload_failures
|
||
|
||
for element in elements[5:]:
|
||
logger.log_queue.append(element)
|
||
|
||
put.failing = False
|
||
await logger.flush_queue()
|
||
|
||
assert logger.log_queue == []
|
||
landed: Final = frozenset(call[0].rsplit("/", 1)[-1] for call in put.calls[-8:])
|
||
assert landed == frozenset(f"test-{index}.json" for index in range(8)) # calls are (url, data) pairs
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.parametrize(
|
||
("status", "code"),
|
||
[
|
||
pytest.param(404, "NoSuchKey", id="404"),
|
||
pytest.param(401, None, id="401"),
|
||
pytest.param(400, None, id="uncoded-400"),
|
||
],
|
||
)
|
||
async def test_unlisted_status_gets_one_put_and_stays_queued(status: int, code: str | None) -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
|
||
put = _FailUntilClearedPut(status=status, code=code)
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
logger.log_queue = [_element({"i": index}, f"{index}") for index in range(4)]
|
||
|
||
with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep:
|
||
await logger.flush_queue()
|
||
|
||
assert len(put.calls) == 4
|
||
mock_sleep.assert_not_awaited()
|
||
assert len(logger.log_queue) == 4
|
||
|
||
|
||
class _SyncRecordingClient:
|
||
def __init__(self, response: httpx.Response) -> None:
|
||
self.response: Final = response
|
||
self.put_calls: list = [] # mutable-ok: call log appended once per PUT
|
||
|
||
def put(
|
||
self,
|
||
url: str,
|
||
data: str | None = None,
|
||
headers: dict[str, str] | None = None,
|
||
timeout: httpx.Timeout | None = None,
|
||
) -> MagicMock:
|
||
self.put_calls.append(url)
|
||
return self.response
|
||
|
||
|
||
def test_sync_upload_404_is_single_attempt_without_sleep() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
|
||
sync_client: Final = _SyncRecordingClient(_coded_failure_response(404, "NoSuchKey"))
|
||
|
||
with (
|
||
patch("litellm.integrations.s3_v2.get_httpx_client", return_value=sync_client),
|
||
patch("time.sleep") as mock_sleep,
|
||
):
|
||
logger.upload_data_to_s3(_element({"id": "sync-404"}, "sync-404"))
|
||
|
||
assert len(sync_client.put_calls) == 1
|
||
mock_sleep.assert_not_called()
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
("status", "expected_puts", "expected_retries"),
|
||
[
|
||
pytest.param(408, 1, 0, id="408-single"),
|
||
pytest.param(502, 1, 0, id="502-single"),
|
||
pytest.param(504, 1, 0, id="504-single"),
|
||
pytest.param(429, 3, 2, id="429-backoff"),
|
||
pytest.param(503, 3, 2, id="503-backoff"),
|
||
],
|
||
)
|
||
def test_sync_upload_retry_set(status: int, expected_puts: int, expected_retries: int) -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
|
||
sync_client: Final = _SyncRecordingClient(_coded_failure_response(status, "SlowDown"))
|
||
|
||
with (
|
||
patch("litellm.integrations.s3_v2.get_httpx_client", return_value=sync_client),
|
||
patch("time.sleep") as mock_sleep,
|
||
):
|
||
logger.upload_data_to_s3(_element({"id": "sync"}, "sync"))
|
||
|
||
assert len(sync_client.put_calls) == expected_puts
|
||
_assert_jittered_backoff(mock_sleep, expected_retries)
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
"failure",
|
||
[pytest.param(httpx.ConnectError, id="connect-error"), pytest.param(httpx.ReadTimeout, id="read-timeout")],
|
||
)
|
||
def test_sync_upload_retries_a_failed_attempt_through_production_http_handler(
|
||
failure: type[httpx.TransportError], caplog: pytest.LogCaptureFixture
|
||
) -> None:
|
||
"""HTTPHandler.put turns an httpx timeout into litellm.Timeout, and each attempt carries its own deadline."""
|
||
requests: Final[list[httpx.Request]] = [] # mutable-ok: the transport records every attempt
|
||
|
||
def respond(request: httpx.Request) -> httpx.Response:
|
||
requests.append(request)
|
||
if len(requests) == 1:
|
||
raise failure("S3 attempt failed", request=request)
|
||
return httpx.Response(200, request=request)
|
||
|
||
handler: Final = HTTPHandler()
|
||
handler.client = httpx.Client(transport=httpx.MockTransport(respond), timeout=None)
|
||
logger: Final = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
with (
|
||
patch( # test-quality-ok: sync upload builds its HTTPHandler per call, there is no injection seam for it
|
||
"litellm.integrations.s3_v2.get_httpx_client", return_value=handler
|
||
),
|
||
patch("time.sleep") as mock_sleep,
|
||
):
|
||
logger.upload_data_to_s3(_element({"id": "sync-transient"}, "sync-transient"))
|
||
handler.client.close()
|
||
|
||
assert len(requests) == 2
|
||
assert requests[0].content == requests[1].content
|
||
assert all(
|
||
request.extensions["timeout"] == httpx.Timeout(S3_UPLOAD_ATTEMPT_TIMEOUT_SECONDS).as_dict()
|
||
for request in requests
|
||
)
|
||
_assert_jittered_backoff(mock_sleep, 1)
|
||
assert "Error uploading to s3" not in caplog.text
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.parametrize(
|
||
("status", "code"),
|
||
[
|
||
pytest.param(503, "SlowDown", id="503"),
|
||
pytest.param(429, "SlowDown", id="429"),
|
||
pytest.param(500, "InternalError", id="500"),
|
||
pytest.param(403, "AccessDenied", id="access-denied-403"),
|
||
],
|
||
)
|
||
async def test_retryable_statuses_back_off_three_attempts(status: int, code: str | None) -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
|
||
put = _FailUntilClearedPut(status=status, code=code)
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
logger.log_queue = [_element({"id": "req"}, "req")]
|
||
|
||
with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep:
|
||
await logger.flush_queue()
|
||
|
||
assert len(put.calls) == 3
|
||
_assert_jittered_backoff(mock_sleep, 2)
|
||
assert len(logger.log_queue) == 1
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.parametrize(
|
||
("status", "code"),
|
||
[
|
||
pytest.param(408, None, id="408"),
|
||
pytest.param(502, None, id="502"),
|
||
pytest.param(504, None, id="504"),
|
||
],
|
||
)
|
||
async def test_unlisted_statuses_are_not_retried_in_call(status: int, code: str | None) -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
|
||
put = _FailUntilClearedPut(status=status, code=code)
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
logger.log_queue = [_element({"id": "req"}, "req")]
|
||
|
||
with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep:
|
||
await logger.flush_queue()
|
||
|
||
assert len(put.calls) == 1
|
||
assert mock_sleep.await_args_list == []
|
||
assert len(logger.log_queue) == 1
|
||
|
||
|
||
class _FirstFailThenOkPut:
|
||
def __init__(self, fail_suffix: str) -> None:
|
||
self.fail_suffix = fail_suffix
|
||
self.failed_once = False
|
||
self.calls: tuple[str, ...] = ()
|
||
|
||
async def __call__(
|
||
self,
|
||
url: str,
|
||
data: str | None = None,
|
||
headers: dict[str, str] | None = None,
|
||
timeout: httpx.Timeout | None = None,
|
||
) -> MagicMock:
|
||
self.calls = (*self.calls, url)
|
||
if url.endswith(self.fail_suffix) and not self.failed_once:
|
||
self.failed_once = True
|
||
return _transient_failure_response()
|
||
return _ok_response()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_retry_finishes_before_the_next_first_attempt() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_max_concurrent_uploads=1,
|
||
)
|
||
|
||
put = _FirstFailThenOkPut("test-a.json")
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
logger.log_queue = [_element({"id": "a"}, "a"), _element({"id": "b"}, "b")]
|
||
|
||
with patch("asyncio.sleep", new=AsyncMock(side_effect=lambda delay: _real_sleep(0))):
|
||
await logger.flush_queue()
|
||
|
||
assert [call_url.rsplit("/", 1)[-1] for call_url in put.calls] == ["test-a.json", "test-a.json", "test-b.json"]
|
||
assert logger.log_queue == []
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_objects_in_backoff_are_bounded_by_the_slot_width() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_max_concurrent_uploads=2,
|
||
)
|
||
|
||
put = _FailUntilClearedPut(status=503)
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
logger.log_queue = [_element({"i": i}, f"{i}") for i in range(20)]
|
||
|
||
sleeping: Final[list[int]] = [0]
|
||
peak: Final[list[int]] = [0]
|
||
|
||
async def counting_sleep(delay: float) -> None:
|
||
sleeping[0] += 1
|
||
peak[0] = max(peak[0], sleeping[0])
|
||
for _ in range(10):
|
||
await _real_sleep(0)
|
||
sleeping[0] -= 1
|
||
|
||
with patch("asyncio.sleep", new=counting_sleep):
|
||
await logger.flush_queue()
|
||
|
||
assert peak[0] <= 2, f"{peak[0]} objects slept at once, slot width is 2"
|
||
assert len(put.calls) == 60
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_subclass_returning_true_drains_the_queue() -> None:
|
||
class _TrueUploadLogger(S3Logger):
|
||
async def async_upload_data_to_s3(self, batch_logging_element: s3BatchLoggingElement) -> bool:
|
||
return True
|
||
|
||
logger = _TrueUploadLogger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.log_queue = [_element({"id": "a"}, "a")]
|
||
|
||
await logger.flush_queue()
|
||
|
||
assert logger.log_queue == []
|
||
logger.async_httpx_client.put.assert_not_called()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_failed_direct_upload_returns_false() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
|
||
put = _FailUntilClearedPut(status=500)
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
test_element = _element({"id": "x"}, "x")
|
||
|
||
with patch("asyncio.sleep", new=AsyncMock(side_effect=lambda delay: _real_sleep(0))):
|
||
outcome = await logger.async_upload_data_to_s3(test_element)
|
||
|
||
assert outcome is False
|
||
assert len(put.calls) == 3
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_terminal_drop_of_one_element_does_not_drop_a_sibling_with_the_same_key() -> None:
|
||
class _TerminalForMarkerPut:
|
||
def __init__(self) -> None:
|
||
self.calls: tuple[str | None, ...] = ()
|
||
|
||
async def __call__(
|
||
self,
|
||
url: str,
|
||
data: str | None = None,
|
||
headers: dict[str, str] | None = None,
|
||
timeout: httpx.Timeout | None = None,
|
||
) -> MagicMock:
|
||
self.calls = (*self.calls, data)
|
||
if data is not None and "terminal-marker" in data:
|
||
return _terminal_failure_response()
|
||
return _transient_failure_response()
|
||
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_drop_on_terminal_error=True,
|
||
)
|
||
|
||
put = _TerminalForMarkerPut()
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
|
||
shared_key = "2025-09-14/shared.json"
|
||
dropped = s3BatchLoggingElement(
|
||
s3_object_key=shared_key, payload={"m": "terminal-marker"}, s3_object_download_filename="shared.json"
|
||
)
|
||
sibling = s3BatchLoggingElement(
|
||
s3_object_key=shared_key, payload={"m": "healthy"}, s3_object_download_filename="shared.json"
|
||
)
|
||
|
||
with patch("asyncio.sleep", new=AsyncMock(side_effect=lambda delay: _real_sleep(0))):
|
||
assert await logger._upload_outcome(dropped) == "dropped"
|
||
assert await logger._upload_outcome(sibling) == "retry"
|
||
|
||
|
||
def test_upload_semaphore_alias_is_the_limiter() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
|
||
assert logger._upload_semaphore is logger._upload_limiter
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_overridden_upload_stays_bounded_by_the_configured_width() -> None:
|
||
class _InFlightUploadLogger(S3Logger):
|
||
def __init__(self, **kwargs: object) -> None:
|
||
super().__init__(**kwargs)
|
||
self.in_flight = 0
|
||
self.peak = 0
|
||
|
||
async def async_upload_data_to_s3(self, batch_logging_element: s3BatchLoggingElement) -> bool:
|
||
self.in_flight += 1
|
||
self.peak = max(self.peak, self.in_flight)
|
||
for _ in range(10):
|
||
await _real_sleep(0)
|
||
self.in_flight -= 1
|
||
return True
|
||
|
||
logger = _InFlightUploadLogger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_max_concurrent_uploads=4,
|
||
)
|
||
logger.log_queue = [_element({"i": index}, f"{index}") for index in range(40)]
|
||
|
||
await logger.flush_queue()
|
||
|
||
assert logger.peak <= 4
|
||
assert logger.log_queue == []
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_holding_the_semaphore_during_a_direct_upload_does_not_deadlock() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_max_concurrent_uploads=1,
|
||
)
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = _RecordingPut()
|
||
element = _element({"id": "x"}, "x")
|
||
|
||
async def held_upload() -> bool:
|
||
async with logger._upload_semaphore:
|
||
return await logger.async_upload_data_to_s3(element)
|
||
|
||
assert await asyncio.wait_for(held_upload(), timeout=5) is True
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_assigning_a_semaphore_changes_the_upload_width() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
logger._upload_semaphore = asyncio.Semaphore(3)
|
||
|
||
put = _CountingPut(width=3)
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
logger.log_queue = [_element({"i": index}, f"{index}") for index in range(30)]
|
||
|
||
await logger.flush_queue()
|
||
|
||
assert put.peak == 3
|
||
assert logger.log_queue == []
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_terminal_code_is_retried_like_base_when_the_drop_flag_is_off() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_drop_on_terminal_error=False,
|
||
)
|
||
put = _StatusPut([_coded_failure_response(403, "InvalidRequest"), _ok_response()])
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
failures = AsyncMock()
|
||
logger.handle_callback_failure = failures
|
||
|
||
with patch("asyncio.sleep", new=AsyncMock(side_effect=lambda delay: _real_sleep(0))):
|
||
assert await logger.async_upload_data_to_s3(_element({"id": "x"}, "x")) is True
|
||
|
||
assert put.calls == 2
|
||
failures.assert_not_called()
|
||
|
||
dropping = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
put.calls = 0
|
||
dropping.async_httpx_client = AsyncMock()
|
||
dropping.async_httpx_client.put = put
|
||
|
||
with patch("asyncio.sleep", new=AsyncMock(side_effect=lambda delay: _real_sleep(0))):
|
||
assert await dropping.async_upload_data_to_s3(_element({"id": "x"}, "x")) is False
|
||
|
||
assert put.calls == 1
|
||
|
||
|
||
def test_sync_terminal_code_is_retried_like_base_when_the_drop_flag_is_off() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_drop_on_terminal_error=False,
|
||
)
|
||
mock_sync_client = MagicMock()
|
||
mock_sync_client.put = MagicMock(side_effect=[_coded_failure_response(403, "InvalidRequest"), _ok_response()])
|
||
failures = MagicMock()
|
||
logger.handle_callback_failure = failures
|
||
|
||
with (
|
||
patch("litellm.integrations.s3_v2.get_httpx_client", return_value=mock_sync_client),
|
||
patch("time.sleep"),
|
||
):
|
||
logger.upload_data_to_s3(_element({"id": "x"}, "x"))
|
||
|
||
assert mock_sync_client.put.call_count == 2
|
||
failures.assert_not_called()
|
||
|
||
dropping = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
mock_sync_client.put = MagicMock(side_effect=[_coded_failure_response(403, "InvalidRequest"), _ok_response()])
|
||
|
||
with (
|
||
patch("litellm.integrations.s3_v2.get_httpx_client", return_value=mock_sync_client),
|
||
patch("time.sleep"),
|
||
):
|
||
dropping.upload_data_to_s3(_element({"id": "x"}, "x"))
|
||
|
||
assert mock_sync_client.put.call_count == 1
|
||
|
||
|
||
def test_sync_retry_lines_stay_at_warning_level(caplog) -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
mock_sync_client = MagicMock()
|
||
mock_sync_client.put = MagicMock(
|
||
side_effect=[_transient_failure_response(503), _transient_failure_response(503), _ok_response()]
|
||
)
|
||
|
||
with (
|
||
caplog.at_level("WARNING"),
|
||
patch("litellm.integrations.s3_v2.get_httpx_client", return_value=mock_sync_client),
|
||
patch("time.sleep"),
|
||
):
|
||
logger.upload_data_to_s3(_element({"id": "x"}, "x"))
|
||
|
||
assert mock_sync_client.put.call_count == 3
|
||
assert sum(1 for record in caplog.records if "retrying in" in record.getMessage()) == 2
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_direct_async_upload_logs_retry_lines_at_warning_level(caplog) -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
put = _StatusPut([_transient_failure_response(503), _ok_response()])
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
|
||
with caplog.at_level("WARNING"), patch("asyncio.sleep", new=AsyncMock(side_effect=lambda delay: _real_sleep(0))):
|
||
assert await logger.async_upload_data_to_s3(_element({"id": "x"}, "x")) is True
|
||
|
||
assert put.calls == 2
|
||
assert sum(1 for record in caplog.records if "retrying in" in record.getMessage()) == 1
|
||
|
||
|
||
def _init_bypassed_logger() -> S3Logger:
|
||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||
|
||
logger = S3Logger.__new__(S3Logger)
|
||
logger.iam_cache = BaseAWSLLM._shared_iam_cache
|
||
logger.s3_endpoint_url = None
|
||
logger.s3_bucket_name = "test-bucket"
|
||
logger.s3_region_name = "us-east-1"
|
||
logger.s3_use_virtual_hosted_style = False
|
||
logger.s3_verify = None
|
||
logger.s3_aws_access_key_id = "test-key"
|
||
logger.s3_aws_secret_access_key = "test-secret"
|
||
logger.s3_aws_session_token = None
|
||
logger.s3_aws_session_name = None
|
||
logger.s3_aws_profile_name = None
|
||
logger.s3_aws_role_name = None
|
||
logger.s3_aws_web_identity_token = None
|
||
logger.s3_aws_sts_endpoint = None
|
||
logger.s3_server_side_encryption = None
|
||
logger.s3_sse_kms_key_id = None
|
||
logger.s3_log_prompts_only = None
|
||
return logger
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_init_bypassed_logger_retries_a_503_and_reports_a_404() -> None:
|
||
logger = _init_bypassed_logger()
|
||
put = _StatusPut([_transient_failure_response(503), _ok_response()])
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
|
||
with patch("asyncio.sleep", new=AsyncMock(side_effect=lambda delay: _real_sleep(0))):
|
||
assert await logger.async_upload_data_to_s3(_element({"id": "x"}, "x")) is True
|
||
|
||
assert put.calls == 2
|
||
|
||
put.calls = 0
|
||
put.responses = [_coded_failure_response(404, None)]
|
||
with patch("asyncio.sleep", new=AsyncMock(side_effect=lambda delay: _real_sleep(0))):
|
||
assert await logger.async_upload_data_to_s3(_element({"id": "y"}, "y")) is False
|
||
|
||
assert put.calls == 1
|
||
|
||
|
||
def test_init_bypassed_sync_logger_retries_a_503_and_reports_a_404() -> None:
|
||
logger = _init_bypassed_logger()
|
||
mock_sync_client = MagicMock()
|
||
mock_sync_client.put = MagicMock(side_effect=[_transient_failure_response(503), _ok_response()])
|
||
|
||
with (
|
||
patch("litellm.integrations.s3_v2.get_httpx_client", return_value=mock_sync_client),
|
||
patch("time.sleep"),
|
||
):
|
||
logger.upload_data_to_s3(_element({"id": "x"}, "x"))
|
||
|
||
assert mock_sync_client.put.call_count == 2
|
||
retried_headers: Final = dict(mock_sync_client.put.call_args.kwargs["headers"])
|
||
assert "X-Amz-Date" in retried_headers
|
||
|
||
mock_sync_client.put = MagicMock(return_value=_coded_failure_response(404, None))
|
||
with (
|
||
patch("litellm.integrations.s3_v2.get_httpx_client", return_value=mock_sync_client),
|
||
patch("time.sleep"),
|
||
):
|
||
logger.upload_data_to_s3(_element({"id": "y"}, "y"))
|
||
|
||
assert mock_sync_client.put.call_count == 1
|
||
failed_url: Final = str(mock_sync_client.put.call_args[0][0])
|
||
assert "test-y.json" in failed_url
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_subclass_with_base_style_upload_bounded_drains_the_queue() -> None:
|
||
class _BaseStyleLogger(S3Logger):
|
||
async def _upload_bounded(self, element: s3BatchLoggingElement) -> bool:
|
||
return True
|
||
|
||
logger = _BaseStyleLogger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.log_queue = [_element({"id": "a"}, "a")]
|
||
|
||
await logger.flush_queue()
|
||
|
||
assert logger.log_queue == []
|
||
logger.async_httpx_client.put.assert_not_called()
|
||
|
||
|
||
def test_bool_config_values_fall_back_to_the_default() -> None:
|
||
from litellm.integrations.s3 import (
|
||
resolve_s3_max_concurrent_uploads,
|
||
resolve_s3_max_queue_size,
|
||
resolve_s3_max_retry_age_seconds,
|
||
)
|
||
|
||
assert resolve_s3_max_concurrent_uploads(True, 16) == 1
|
||
assert resolve_s3_max_queue_size(True, 50000) == 50000
|
||
assert resolve_s3_max_retry_age_seconds(True, 3600) == 3600
|
||
|
||
|
||
def test_int_env_helper_falls_back_on_non_numeric(monkeypatch: pytest.MonkeyPatch) -> None:
|
||
from litellm.litellm_core_utils.env_utils import get_env_int
|
||
|
||
monkeypatch.setenv("TEST_S3_INT_ENV", "abc")
|
||
assert get_env_int("TEST_S3_INT_ENV", 3) == 3
|
||
monkeypatch.setenv("TEST_S3_INT_ENV", "7")
|
||
assert get_env_int("TEST_S3_INT_ENV", 3) == 7
|
||
|
||
|
||
class _FailOncePerKeyPut:
|
||
def __init__(self) -> None:
|
||
self.failed: set[str] = set()
|
||
self.calls: tuple[str, ...] = ()
|
||
|
||
async def __call__(
|
||
self,
|
||
url: str,
|
||
data: str | None = None,
|
||
headers: dict[str, str] | None = None,
|
||
timeout: httpx.Timeout | None = None,
|
||
) -> MagicMock:
|
||
self.calls = (*self.calls, url)
|
||
if url not in self.failed:
|
||
self.failed.add(url)
|
||
return _transient_failure_response()
|
||
return _ok_response()
|
||
|
||
|
||
class _SlowFailOncePerKeyPut:
|
||
def __init__(self, dumps_count) -> None:
|
||
self.failed: set[str] = set()
|
||
self.dumps_count = dumps_count
|
||
self.first_completed: int | None = None
|
||
self.calls: tuple[str, ...] = ()
|
||
|
||
async def __call__(
|
||
self,
|
||
url: str,
|
||
data: str | None = None,
|
||
headers: dict[str, str] | None = None,
|
||
timeout: httpx.Timeout | None = None,
|
||
) -> MagicMock:
|
||
self.calls = (*self.calls, url)
|
||
await _real_sleep(0)
|
||
if self.first_completed is None:
|
||
self.first_completed = self.dumps_count()
|
||
if url not in self.failed:
|
||
self.failed.add(url)
|
||
return _transient_failure_response()
|
||
return _ok_response()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_peak_serialized_bodies_bounded_by_upload_width() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
|
||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps as real_safe_dumps
|
||
|
||
dumps_calls: list[object] = []
|
||
|
||
def counting_dumps(*args, **kwargs):
|
||
dumps_calls.append(args)
|
||
return real_safe_dumps(*args, **kwargs)
|
||
|
||
put = _SlowFailOncePerKeyPut(lambda: len(dumps_calls))
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
logger.log_queue = [_element({"i": index}, f"{index}") for index in range(64)]
|
||
|
||
with (
|
||
patch("litellm.integrations.s3_v2.safe_dumps", side_effect=counting_dumps),
|
||
patch("asyncio.sleep", new=AsyncMock(side_effect=lambda delay: _real_sleep(0))),
|
||
):
|
||
await logger.flush_queue()
|
||
|
||
assert put.first_completed is not None
|
||
assert put.first_completed <= logger.s3_max_concurrent_uploads
|
||
assert len(dumps_calls) == 64
|
||
assert len(put.calls) == 128
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_send_batch_calls_upload_with_one_positional_arg() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
|
||
uploaded: list[str] = [] # mutable-ok: appended once per upload by the double
|
||
|
||
async def mock_upload(batch_logging_element) -> str:
|
||
uploaded.append(batch_logging_element.s3_object_key)
|
||
return "delivered"
|
||
|
||
logger.async_upload_data_to_s3 = mock_upload
|
||
logger.log_queue = [_element({"id": "a"}, "a"), _element({"id": "b"}, "b")]
|
||
|
||
await logger.flush_queue()
|
||
|
||
assert sorted(key.rsplit("/", 1)[-1] for key in uploaded) == ["test-a.json", "test-b.json"]
|
||
assert logger.log_queue == []
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_retries_serialize_the_body_once_per_element_per_flush() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
|
||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps as real_safe_dumps
|
||
|
||
dumps_calls: list[object] = []
|
||
|
||
def counting_dumps(*args, **kwargs):
|
||
dumps_calls.append(args)
|
||
return real_safe_dumps(*args, **kwargs)
|
||
|
||
put = _FailUntilClearedPut(status=503)
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
logger.log_queue = [_element({"i": index}, f"{index}") for index in range(8)]
|
||
|
||
with (
|
||
patch("litellm.integrations.s3_v2.safe_dumps", side_effect=counting_dumps),
|
||
patch("asyncio.sleep", new=AsyncMock(side_effect=lambda delay: _real_sleep(0))),
|
||
):
|
||
await logger.flush_queue()
|
||
|
||
assert len(dumps_calls) == 8
|
||
assert len(put.calls) == 24
|
||
assert len(logger.log_queue) == 8
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_flush_logs_one_retry_warning(caplog) -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
|
||
put = _FailOncePerKeyPut()
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
logger.log_queue = [_element({"i": index}, f"{index}") for index in range(8)]
|
||
|
||
with caplog.at_level("WARNING"), patch("asyncio.sleep", new_callable=AsyncMock):
|
||
await logger.flush_queue()
|
||
|
||
assert logger.log_queue == []
|
||
assert sum(1 for record in caplog.records if "in-call retries" in record.getMessage()) == 1
|
||
assert all("retrying in" not in record.getMessage() for record in caplog.records)
|
||
|
||
|
||
class _AppendingSuffixFailingPut:
|
||
def __init__(self, logger: S3Logger, element: s3BatchLoggingElement, fail_suffixes: tuple[str, ...]) -> None:
|
||
self.logger = logger
|
||
self.element = element
|
||
self.fail_suffixes = fail_suffixes
|
||
self.appended = False
|
||
|
||
async def __call__(
|
||
self,
|
||
url: str,
|
||
data: str | None = None,
|
||
headers: dict[str, str] | None = None,
|
||
timeout: httpx.Timeout | None = None,
|
||
) -> MagicMock:
|
||
if not self.appended:
|
||
self.appended = True
|
||
self.logger.log_queue.append(self.element)
|
||
if url.endswith(self.fail_suffixes):
|
||
return _transient_failure_response()
|
||
return _ok_response()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_failed_elements_stay_oldest_first_when_requeued() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_max_retry_age_seconds=3600,
|
||
)
|
||
|
||
late = _element({"id": "late"}, "late")
|
||
failed = _element({"id": "f2"}, "f2")
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = _AppendingSuffixFailingPut(logger, late, ("test-f2.json",))
|
||
logger.log_queue = [_element({"id": "f1"}, "f1"), failed]
|
||
|
||
with patch("asyncio.sleep", new_callable=AsyncMock):
|
||
await logger.flush_queue()
|
||
|
||
assert [element.s3_object_key for element in logger.log_queue] == [failed.s3_object_key, late.s3_object_key]
|
||
assert logger.log_queue[0].retrying_since is not None
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_overflow_prefers_arrivals_over_failed_elements_without_counting_the_trim(caplog) -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_max_queue_size=1,
|
||
)
|
||
|
||
late = _element({"id": "late"}, "late")
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = _AppendingSuffixFailingPut(logger, late, ("test-f2.json",))
|
||
logger.log_queue = [_element({"id": "f1"}, "f1"), _element({"id": "f2"}, "f2")]
|
||
|
||
with (
|
||
patch.object(logger, "handle_callback_failure") as mock_failure,
|
||
patch("asyncio.sleep", new_callable=AsyncMock),
|
||
):
|
||
await logger.flush_queue()
|
||
|
||
assert [element.s3_object_key for element in logger.log_queue] == [late.s3_object_key]
|
||
failed_uploads: Final = 1
|
||
assert mock_failure.call_count == failed_uploads
|
||
assert "dropped 1 oldest events" in caplog.text
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_fresh_elements_upload_before_stale_retries_after_a_failed_flush() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_max_concurrent_uploads=1,
|
||
)
|
||
|
||
late = _element({"id": "late"}, "late")
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = _AppendingSuffixFailingPut(logger, late, ("test-f2.json",))
|
||
logger.log_queue = [_element({"id": "f1"}, "f1"), _element({"id": "f2"}, "f2")]
|
||
|
||
with patch("asyncio.sleep", new_callable=AsyncMock):
|
||
await logger.flush_queue()
|
||
|
||
recovered = _FailOnSuffixPut(("never-matches",))
|
||
logger.async_httpx_client.put = recovered
|
||
with patch("asyncio.sleep", new_callable=AsyncMock):
|
||
await logger.flush_queue()
|
||
|
||
assert [call_url.rsplit("/", 1)[-1] for call_url in recovered.calls] == ["test-late.json", "test-f2.json"]
|
||
assert logger.log_queue == []
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_repeated_overflow_trims_oldest_across_failed_flushes(caplog) -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_max_queue_size=3,
|
||
)
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = _AppendingFailingPut(logger, (_element({"id": "d"}, "d"),))
|
||
logger.log_queue = [_element({"id": name}, name) for name in ("a", "b", "c")]
|
||
|
||
with patch("asyncio.sleep", new_callable=AsyncMock):
|
||
await logger.flush_queue()
|
||
|
||
assert [element.payload["id"] for element in logger.log_queue] == ["b", "c", "d"]
|
||
|
||
logger.async_httpx_client.put = _AppendingFailingPut(logger, (_element({"id": "e"}, "e"),))
|
||
with patch("asyncio.sleep", new_callable=AsyncMock):
|
||
await logger.flush_queue()
|
||
|
||
assert [element.payload["id"] for element in logger.log_queue] == ["c", "d", "e"]
|
||
assert caplog.text.count("dropped 1 oldest events") == 2
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_retry_age_budget_drops_after_the_clock_set_by_a_partial_failure(caplog) -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_max_retry_age_seconds=1,
|
||
)
|
||
|
||
put = _FailOnSuffixCodedPut(("test-poison.json",), 503, "SlowDown")
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
logger.log_queue = [_element({"id": "poison"}, "poison"), _element({"id": "good"}, "good")]
|
||
|
||
t0: Final = _NOW
|
||
with patch.object(logger, "handle_callback_failure") as mock_failure:
|
||
with (
|
||
patch("asyncio.sleep", new_callable=AsyncMock),
|
||
patch("time.monotonic", return_value=t0),
|
||
):
|
||
await logger.flush_queue()
|
||
|
||
assert [element.s3_object_key for element in logger.log_queue] == ["2025-09-14/test-poison.json"]
|
||
assert logger.log_queue[0].retrying_since == t0
|
||
|
||
logger.log_queue.append(_element({"id": "good-2"}, "good-2"))
|
||
with (
|
||
patch("asyncio.sleep", new_callable=AsyncMock),
|
||
patch("time.monotonic", return_value=t0 + 2),
|
||
):
|
||
await logger.flush_queue()
|
||
|
||
assert logger.log_queue == []
|
||
assert "retrying longer than s3_max_retry_age_seconds=1" in caplog.text
|
||
poison_puts: Final = sum(1 for call_url in put.calls if call_url.endswith("test-poison.json"))
|
||
assert poison_puts == 6
|
||
upload_failures: Final = 2
|
||
assert mock_failure.call_count == upload_failures
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_the_retry_clock_starts_at_the_first_partial_failure_not_first_seen() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_max_retry_age_seconds=1,
|
||
)
|
||
|
||
put = _FailUntilClearedPut(status=503, code="SlowDown")
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
logger.log_queue = [_element({"id": "poison"}, "poison")]
|
||
|
||
t0: Final = _NOW
|
||
with (
|
||
patch("asyncio.sleep", new_callable=AsyncMock),
|
||
patch("time.monotonic", return_value=t0),
|
||
):
|
||
await logger.flush_queue()
|
||
|
||
assert len(logger.log_queue) == 1
|
||
assert logger.log_queue[0].retrying_since is None
|
||
|
||
logger.log_queue.append(_element({"id": "good"}, "good"))
|
||
put.failing = False
|
||
failing_poison: Final = _FailOnSuffixCodedPut(("test-poison.json",), 503, "SlowDown")
|
||
logger.async_httpx_client.put = failing_poison
|
||
with (
|
||
patch("asyncio.sleep", new_callable=AsyncMock),
|
||
patch("time.monotonic", return_value=t0 + 500),
|
||
):
|
||
await logger.flush_queue()
|
||
|
||
assert [element.s3_object_key for element in logger.log_queue] == ["2025-09-14/test-poison.json"]
|
||
assert logger.log_queue[0].retrying_since == t0 + 500
|
||
|
||
|
||
def test_sync_upload_retries_access_denied_403(caplog):
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
|
||
test_element = s3BatchLoggingElement(
|
||
s3_object_key="2025-09-14/test-sync-403.json",
|
||
payload={"test": "sync-403"},
|
||
s3_object_download_filename="test-sync-403.json",
|
||
)
|
||
|
||
mock_sync_client = MagicMock()
|
||
mock_sync_client.put = MagicMock(return_value=_coded_failure_response(403, "AccessDenied"))
|
||
|
||
with (
|
||
patch("litellm.integrations.s3_v2.get_httpx_client", return_value=mock_sync_client),
|
||
patch("time.sleep") as mock_sleep,
|
||
):
|
||
logger.upload_data_to_s3(test_element)
|
||
|
||
assert mock_sync_client.put.call_count == 3
|
||
_assert_jittered_backoff(mock_sleep, 2)
|
||
assert "dropping object" not in caplog.text
|
||
|
||
|
||
def test_sync_upload_drops_terminal_object_once_and_logs_it_only_when_opted_in(caplog):
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_drop_on_terminal_error=True,
|
||
)
|
||
|
||
test_element = s3BatchLoggingElement(
|
||
s3_object_key="2025-09-14/test-sync-terminal.json",
|
||
payload={"test": "sync-terminal"},
|
||
s3_object_download_filename="test-sync-terminal.json",
|
||
)
|
||
|
||
mock_sync_client = MagicMock()
|
||
mock_sync_client.put = MagicMock(return_value=_coded_failure_response(400, "EntityTooLarge"))
|
||
|
||
with (
|
||
patch("litellm.integrations.s3_v2.get_httpx_client", return_value=mock_sync_client),
|
||
patch("time.sleep") as mock_sleep,
|
||
):
|
||
logger.upload_data_to_s3(test_element)
|
||
|
||
assert mock_sync_client.put.call_count == 1
|
||
mock_sleep.assert_not_called()
|
||
assert "dropping object" in caplog.text
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_requeued_batch_file_keeps_the_earliest_member_retrying_since() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_batch_file_upload=True,
|
||
)
|
||
|
||
put = _FailUntilClearedPut()
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
|
||
stale: Final = _NOW - 30
|
||
retried = _element({"id": "retried"}, "retried").model_copy(update={"retrying_since": stale})
|
||
fresh = _element({"id": "fresh"}, "fresh")
|
||
logger.log_queue = [retried, fresh]
|
||
|
||
with (
|
||
patch("asyncio.sleep", new_callable=AsyncMock),
|
||
patch("time.monotonic", return_value=_NOW),
|
||
):
|
||
await logger.flush_queue()
|
||
|
||
assert len(logger.log_queue) == 1
|
||
assert logger.log_queue[0].s3_object_key.endswith(".jsonl")
|
||
assert logger.log_queue[0].retrying_since == stale
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_an_unlisted_5xx_is_requeued_without_an_extra_attempt() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
|
||
put = _FailUntilClearedPut(status=507)
|
||
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = put
|
||
logger.log_queue = [_element({"id": "req-507"}, "507")]
|
||
|
||
with patch("asyncio.sleep", new_callable=AsyncMock):
|
||
await logger.flush_queue()
|
||
|
||
assert len(logger.log_queue) == 1
|
||
assert len(put.calls) == 1
|
||
|
||
|
||
@pytest.mark.parametrize("configured", [0, "0", None, ""])
|
||
def test_retry_age_resolution_disables_the_budget(configured: object) -> None:
|
||
from litellm.integrations.s3 import resolve_s3_max_retry_age_seconds
|
||
|
||
assert resolve_s3_max_retry_age_seconds(configured, 3600) is None
|
||
|
||
|
||
@pytest.mark.parametrize("configured", ["abc", -5, True])
|
||
def test_invalid_retry_age_resolution_falls_back_with_a_warning(configured: object, caplog) -> None:
|
||
from litellm.integrations.s3 import resolve_s3_max_retry_age_seconds
|
||
|
||
assert resolve_s3_max_retry_age_seconds(configured, 3600) == 3600
|
||
assert "s3_max_retry_age_seconds" in caplog.text
|
||
|
||
|
||
def test_retry_age_resolution_accepts_a_positive_int() -> None:
|
||
from litellm.integrations.s3 import resolve_s3_max_retry_age_seconds
|
||
|
||
assert resolve_s3_max_retry_age_seconds(30, 3600) == 30
|
||
|
||
|
||
def test_default_logger_sets_a_one_hour_retry_age_budget() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
)
|
||
|
||
assert logger.s3_max_retry_age_seconds == 3600
|
||
|
||
|
||
def test_constructor_zero_disables_the_retry_age_budget() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_max_retry_age_seconds=0,
|
||
)
|
||
|
||
assert logger.s3_max_retry_age_seconds is None
|
||
|
||
|
||
def test_invalid_callback_params_retry_age_falls_back_to_the_default() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_callback_params_override={"s3_max_retry_age_seconds": "abc"},
|
||
)
|
||
|
||
assert logger.s3_max_retry_age_seconds == 3600
|
||
|
||
|
||
def test_callback_params_retry_age_wins_over_constructor() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_max_retry_age_seconds=30,
|
||
s3_callback_params_override={"s3_max_retry_age_seconds": 60},
|
||
)
|
||
|
||
assert logger.s3_max_retry_age_seconds == 60
|
||
|
||
|
||
def test_callback_params_drop_terminal_error_wins_over_constructor() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_drop_on_terminal_error=False,
|
||
s3_callback_params_override={"s3_drop_on_terminal_error": True},
|
||
)
|
||
|
||
assert logger.s3_drop_on_terminal_error is True
|
||
|
||
|
||
def test_invalid_callback_params_drop_terminal_error_falls_back_to_constructor_value() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_drop_on_terminal_error=True,
|
||
s3_callback_params_override={"s3_drop_on_terminal_error": "banana"},
|
||
)
|
||
|
||
assert logger.s3_drop_on_terminal_error is True
|
||
|
||
|
||
def test_callback_params_adaptive_concurrency_wins_over_constructor() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_adaptive_concurrency=False,
|
||
s3_callback_params_override={"s3_adaptive_concurrency": "true"},
|
||
)
|
||
|
||
assert logger.s3_adaptive_concurrency is True
|
||
assert logger._upload_limiter._ceiling > logger._upload_limiter.limit
|
||
|
||
|
||
def test_invalid_callback_params_max_adaptive_concurrency_falls_back_to_default() -> None:
|
||
from litellm.constants import DEFAULT_S3_MAX_ADAPTIVE_CONCURRENCY
|
||
|
||
logger = _override_logger(s3_adaptive_concurrency=True, s3_max_adaptive_concurrency="abc")
|
||
|
||
assert logger.s3_max_adaptive_concurrency == DEFAULT_S3_MAX_ADAPTIVE_CONCURRENCY
|
||
assert logger._upload_limiter._ceiling == DEFAULT_S3_MAX_ADAPTIVE_CONCURRENCY
|
||
|
||
|
||
def test_callback_params_queue_size_wins_over_constructor() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_max_queue_size=7,
|
||
s3_callback_params_override={"s3_max_queue_size": 4},
|
||
)
|
||
|
||
assert logger.s3_max_queue_size == 4
|
||
assert logger.max_queue_size == 4
|
||
|
||
|
||
def test_invalid_callback_params_queue_size_falls_back_to_constructor_value() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_max_queue_size=7,
|
||
s3_callback_params_override={"s3_max_queue_size": "abc"},
|
||
)
|
||
|
||
assert logger.s3_max_queue_size == 7
|
||
|
||
|
||
def test_invalid_constructor_queue_size_falls_back_to_default() -> None:
|
||
from litellm.integrations.custom_batch_logger import CustomBatchLogger
|
||
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_max_queue_size="abc",
|
||
)
|
||
|
||
assert logger.s3_max_queue_size == CustomBatchLogger.DEFAULT_MAX_QUEUE_SIZE
|
||
assert logger.max_queue_size == CustomBatchLogger.DEFAULT_MAX_QUEUE_SIZE
|
||
|
||
|
||
class _StatusPut:
|
||
def __init__(self, responses: "list[MagicMock | Exception]") -> None:
|
||
self.responses = responses
|
||
self.calls = 0
|
||
|
||
async def __call__(
|
||
self,
|
||
url: str,
|
||
data: str | None = None,
|
||
headers: dict[str, str] | None = None,
|
||
timeout: httpx.Timeout | None = None,
|
||
) -> MagicMock:
|
||
self.calls += 1
|
||
outcome = self.responses[min(self.calls - 1, len(self.responses) - 1)]
|
||
if isinstance(outcome, Exception):
|
||
raise outcome
|
||
return outcome
|
||
|
||
|
||
def _slow_down_response(status: int = 200) -> MagicMock:
|
||
response = _ok_response() if status == 200 else _transient_failure_response(status)
|
||
response.text = "<Error><Code>SlowDown</Code></Error>"
|
||
return response
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_503_response_lowers_the_adaptive_limit() -> None:
|
||
logger = _override_logger(s3_adaptive_concurrency=True)
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = _StatusPut([_transient_failure_response(503), _ok_response()])
|
||
|
||
logger._upload_limiter._limit = 64 # mutable-ok: seed the AIMD state above the floor without replaying growth
|
||
with patch("asyncio.sleep", new_callable=AsyncMock):
|
||
logger.log_queue = [_element({"i": 0}, "0")]
|
||
await logger.flush_queue()
|
||
|
||
assert logger._upload_limiter.limit == 32
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_429_response_lowers_the_adaptive_limit() -> None:
|
||
logger = _override_logger(s3_adaptive_concurrency=True)
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = _StatusPut([_transient_failure_response(429), _ok_response()])
|
||
|
||
logger._upload_limiter._limit = 64 # mutable-ok: seed the AIMD state above the floor without replaying growth
|
||
with patch("asyncio.sleep", new_callable=AsyncMock):
|
||
logger.log_queue = [_element({"i": 0}, "0")]
|
||
await logger.flush_queue()
|
||
|
||
assert logger._upload_limiter.limit == 32
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_slow_down_body_code_lowers_the_adaptive_limit() -> None:
|
||
logger = _override_logger(s3_adaptive_concurrency=True)
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = _StatusPut([_slow_down_response()])
|
||
|
||
logger._upload_limiter._limit = 64 # mutable-ok: seed the AIMD state above the floor without replaying growth
|
||
logger.log_queue = [_element({"i": 0}, "0")]
|
||
await logger.async_send_batch()
|
||
|
||
assert logger._upload_limiter.limit == 32
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_transport_error_lowers_the_adaptive_limit() -> None:
|
||
logger = _override_logger(s3_adaptive_concurrency=True)
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = _StatusPut(
|
||
[httpx.ConnectError("connect refused", request=MagicMock()), _ok_response()]
|
||
)
|
||
|
||
logger._upload_limiter._limit = 64 # mutable-ok: seed the AIMD state above the floor without replaying growth
|
||
with patch("asyncio.sleep", new_callable=AsyncMock):
|
||
logger.log_queue = [_element({"i": 0}, "0")]
|
||
await logger.flush_queue()
|
||
|
||
assert logger._upload_limiter.limit == 32
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_fast_uploads_raise_the_adaptive_limit() -> None:
|
||
logger = _override_logger(s3_adaptive_concurrency=True)
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = _RecordingPut()
|
||
|
||
before: Final = logger._upload_limiter.limit
|
||
logger.log_queue = [_element({"i": i}, f"{i}") for i in range(before)]
|
||
await logger.async_send_batch()
|
||
|
||
assert logger._upload_limiter.limit > before
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_configured_concurrency_is_the_fixed_limit_when_adaptive_is_off() -> None:
|
||
logger = S3Logger(
|
||
s3_bucket_name="test-bucket",
|
||
s3_aws_access_key_id="test-key",
|
||
s3_aws_secret_access_key="test-secret",
|
||
s3_region_name="us-east-1",
|
||
s3_max_concurrent_uploads=64,
|
||
)
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = _StatusPut([_transient_failure_response(503), _ok_response()])
|
||
|
||
assert logger._upload_limiter._value == 64
|
||
|
||
logger.log_queue = [_element({"i": 0}, "0")]
|
||
with patch("asyncio.sleep", new_callable=AsyncMock):
|
||
await logger.flush_queue()
|
||
|
||
assert logger._upload_limiter._value == 64
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_the_limit_never_falls_below_the_configured_width() -> None:
|
||
logger = _override_logger(s3_adaptive_concurrency=True, s3_max_concurrent_uploads=8)
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = _StatusPut(
|
||
[_transient_failure_response(503), _transient_failure_response(503), _transient_failure_response(503)]
|
||
)
|
||
|
||
with patch("asyncio.sleep", new_callable=AsyncMock):
|
||
logger.log_queue = [_element({"i": 0}, "0")]
|
||
await logger.flush_queue()
|
||
|
||
assert logger._upload_limiter.limit == 8
|
||
|
||
|
||
def test_default_upload_width_is_16() -> None:
|
||
from litellm.constants import DEFAULT_S3_MAX_CONCURRENT_UPLOADS
|
||
|
||
logger = _override_logger()
|
||
|
||
assert logger._upload_semaphore._value == DEFAULT_S3_MAX_CONCURRENT_UPLOADS
|
||
assert DEFAULT_S3_MAX_CONCURRENT_UPLOADS == 16
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_a_slow_put_does_not_lower_the_adaptive_limit() -> None:
|
||
logger = _override_logger(s3_adaptive_concurrency=True)
|
||
logger.async_httpx_client = AsyncMock()
|
||
|
||
async def slow_put(
|
||
url: str, data: str | None = None, headers: dict[str, str] | None = None, timeout: httpx.Timeout | None = None
|
||
) -> MagicMock:
|
||
await _real_sleep(0)
|
||
return _ok_response()
|
||
|
||
logger.async_httpx_client.put = slow_put
|
||
|
||
before: Final = logger._upload_limiter.limit
|
||
logger.log_queue = [_element({"i": 0}, "0")]
|
||
await logger.async_send_batch()
|
||
|
||
assert logger._upload_limiter.limit >= before
|
||
|
||
|
||
class _FastOkPut:
|
||
def __init__(self) -> None:
|
||
self.calls = 0
|
||
|
||
async def __call__(
|
||
self,
|
||
url: str,
|
||
data: str | None = None,
|
||
headers: dict[str, str] | None = None,
|
||
timeout: httpx.Timeout | None = None,
|
||
) -> MagicMock:
|
||
self.calls += 1
|
||
await _real_sleep(0)
|
||
return _ok_response()
|
||
|
||
|
||
async def _timed_send_batch(size: int) -> float:
|
||
logger = _override_logger()
|
||
logger.async_httpx_client = AsyncMock()
|
||
logger.async_httpx_client.put = _FastOkPut()
|
||
logger.log_queue = [_element({"i": i}, f"{i}") for i in range(size)]
|
||
started = time.perf_counter()
|
||
await logger.async_send_batch()
|
||
return time.perf_counter() - started
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_send_batch_time_grows_linearly_with_the_batch() -> None:
|
||
baseline: Final = await _timed_send_batch(2_000)
|
||
quadrupled: Final = await _timed_send_batch(8_000)
|
||
|
||
assert quadrupled / baseline < 8, f"2k took {baseline:.3f}s, 8k took {quadrupled:.3f}s"
|
||
|
||
|
||
class QueueOnlyS3LoggerWithoutCredentialsOrPeriodicFlush(S3Logger):
|
||
def __init__(self) -> None:
|
||
CustomLogger.__init__(self)
|
||
self.log_queue: list[s3BatchLoggingElement] = []
|
||
self.batch_size = 100
|
||
self.s3_strip_base64_files = False
|
||
self.s3_use_team_prefix = False
|
||
self.s3_use_key_prefix = False
|
||
self.s3_path = ""
|
||
self.s3_log_prompts_only = None
|
||
self.s3_batch_file_upload = False
|
||
self._upload_semaphore = asyncio.Semaphore(1)
|
||
|
||
|
||
@pytest.mark.parametrize("stream", [False, True])
|
||
@pytest.mark.parametrize("concurrent", [False, True])
|
||
@pytest.mark.asyncio
|
||
async def test_repeated_failure_notifications_upload_once(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
stream: bool,
|
||
concurrent: bool,
|
||
) -> None:
|
||
sink: Final = QueueOnlyS3LoggerWithoutCredentialsOrPeriodicFlush()
|
||
upload: Final = AsyncMock(return_value=True)
|
||
monkeypatch.setattr(sink, "async_upload_data_to_s3", upload)
|
||
monkeypatch.setattr(litellm, "callbacks", [])
|
||
monkeypatch.setattr(litellm, "_async_failure_callback", [sink])
|
||
logging_obj: Final = Logging(
|
||
model="openai/gpt-4o-mini",
|
||
messages=[],
|
||
stream=stream,
|
||
call_type="acompletion",
|
||
start_time=datetime.now(),
|
||
litellm_call_id="synthetic-request",
|
||
function_id="synthetic-function",
|
||
)
|
||
error: Final = litellm.ServiceUnavailableError(
|
||
message="synthetic deployment-selection failure",
|
||
llm_provider="openai",
|
||
model="openai/gpt-4o-mini",
|
||
)
|
||
|
||
if concurrent:
|
||
await asyncio.gather(
|
||
*(logging_obj.async_failure_handler(error, "synthetic traceback") for _ in range(3))
|
||
)
|
||
else:
|
||
for _ in range(3):
|
||
await logging_obj.async_failure_handler(error, "synthetic traceback")
|
||
|
||
assert len(sink.log_queue) == 1
|
||
assert all(entry.payload["status"] == "failure" for entry in sink.log_queue)
|
||
await sink.async_send_batch()
|
||
await asyncio.sleep(0)
|
||
assert upload.await_count == 1
|
||
|
||
|
||
@pytest.fixture()
|
||
def _vcr_outcome_gate(request, vcr):
|
||
install_live_call_probe(request, vcr)
|
||
yield
|
||
record_vcr_outcome(request, vcr)
|
||
|
||
@pytest_asyncio.fixture(loop_scope="function")
|
||
async def drain_logging_worker(isolate_litellm_state: None) -> AsyncIterator[None]:
|
||
yield
|
||
await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=LOGGING_WORKER_DRAIN_TIMEOUT_SECONDS)
|
||
|
||
LOGGING_WORKER_DRAIN_TIMEOUT_SECONDS: Final = LOGGING_WORKER_MAX_TIME_PER_COROUTINE + 5.0
|
||
|
||
@pytest.fixture(scope="function", autouse=True)
|
||
def isolate_litellm_state():
|
||
"""
|
||
Per-function isolation fixture.
|
||
|
||
Resets litellm state to the true defaults captured at conftest import time,
|
||
then restores after the test. This prevents module-level mutations (e.g.
|
||
`litellm.num_retries = 3` at the top of test_langfuse_e2e_test.py) from
|
||
leaking across tests within the same xdist worker.
|
||
"""
|
||
from litellm.litellm_core_utils import litellm_logging as ll_logging
|
||
from litellm.proxy.management_helpers import audit_logs as ll_audit_logs
|
||
|
||
if hasattr(litellm, "in_memory_llm_clients_cache"):
|
||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||
ll_logging._in_memory_loggers.clear()
|
||
ll_audit_logs._audit_log_callback_cache.clear()
|
||
for attr in _LIST_ATTRS:
|
||
if attr in _DEFAULTS:
|
||
default = _DEFAULTS[attr]
|
||
setattr(litellm, attr, default.copy() if isinstance(default, list) else default)
|
||
for attr in _SCALAR_ATTRS:
|
||
if attr in _DEFAULTS:
|
||
setattr(litellm, attr, _DEFAULTS[attr])
|
||
yield
|
||
if hasattr(litellm, "in_memory_llm_clients_cache"):
|
||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||
ll_logging._in_memory_loggers.clear()
|
||
ll_audit_logs._audit_log_callback_cache.clear()
|
||
for attr in _LIST_ATTRS:
|
||
if attr in _DEFAULTS:
|
||
default = _DEFAULTS[attr]
|
||
setattr(litellm, attr, default.copy() if isinstance(default, list) else default)
|
||
for attr in _SCALAR_ATTRS:
|
||
if attr in _DEFAULTS:
|
||
setattr(litellm, attr, _DEFAULTS[attr])
|
||
|
||
_LIST_ATTRS = (
|
||
"callbacks",
|
||
"success_callback",
|
||
"failure_callback",
|
||
"_async_success_callback",
|
||
"_async_failure_callback",
|
||
"service_callback",
|
||
"pre_call_rules",
|
||
"post_call_rules",
|
||
)
|
||
|
||
_SCALAR_ATTRS = (
|
||
"set_verbose",
|
||
"cache",
|
||
"num_retries",
|
||
"num_retries_per_request",
|
||
"turn_off_message_logging",
|
||
"redact_messages_in_exceptions",
|
||
"redact_user_api_key_info",
|
||
"s3_callback_params",
|
||
"s3_audit_callback_params",
|
||
"datadog_params",
|
||
"vector_store_registry",
|
||
)
|
||
|
||
_DEFAULTS: Final = {
|
||
attr: getattr(litellm, attr).copy() if isinstance(getattr(litellm, attr), list) else getattr(litellm, attr)
|
||
for attr in (*_LIST_ATTRS, *_SCALAR_ATTRS)
|
||
if hasattr(litellm, attr)
|
||
}
|
||
|
||
@pytest.fixture(scope="module")
|
||
def setup_and_teardown():
|
||
"""
|
||
Module-scoped setup. Reloads litellm only in single-process mode
|
||
(skipped under xdist to avoid cross-worker interference).
|
||
"""
|
||
import litellm
|
||
|
||
worker_id = os.environ.get("PYTEST_XDIST_WORKER", None)
|
||
if worker_id is None:
|
||
importlib.reload(litellm)
|
||
try:
|
||
if hasattr(litellm, "proxy") and hasattr(litellm.proxy, "proxy_server"):
|
||
import litellm.proxy.proxy_server
|
||
|
||
importlib.reload(litellm.proxy.proxy_server)
|
||
except Exception:
|
||
pass
|
||
if hasattr(litellm, "in_memory_llm_clients_cache"):
|
||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||
yield
|
||
|
||
@pytest.fixture
|
||
def amazing_s3_retries(monkeypatch: pytest.MonkeyPatch) -> None:
|
||
monkeypatch.setattr(litellm, "num_retries", 3)
|
||
|
||
class _FakeS3Paginator:
|
||
def __init__(self, objects):
|
||
self.objects = objects
|
||
|
||
def paginate(self, Bucket):
|
||
keys = sorted(self.objects[Bucket])
|
||
if not keys:
|
||
return [{}]
|
||
return [{"Contents": [{"Key": key} for key in keys]}]
|
||
|
||
class _FakeS3Client:
|
||
def __init__(self):
|
||
self.objects = defaultdict(dict)
|
||
self.object_written = threading.Event()
|
||
|
||
def clear(self):
|
||
self.objects.clear()
|
||
self.object_written.clear()
|
||
|
||
def put_object(self, Bucket, Key, Body, **_kwargs):
|
||
self.objects[Bucket][Key] = Body
|
||
self.object_written.set()
|
||
return {"ResponseMetadata": {"HTTPStatusCode": 200}}
|
||
|
||
def delete_object(self, Bucket, Key):
|
||
self.objects[Bucket].pop(Key, None)
|
||
return {"ResponseMetadata": {"HTTPStatusCode": 204}}
|
||
|
||
def get_paginator(self, name):
|
||
assert name == "list_objects_v2"
|
||
return _FakeS3Paginator(self.objects)
|
||
|
||
def list_objects(self, Bucket):
|
||
keys = sorted(self.objects[Bucket])
|
||
return {"Contents": [{"Key": key, "LastModified": 0} for key in keys]}
|
||
|
||
_FAKE_S3_CLIENT = _FakeS3Client()
|
||
|
||
@pytest.fixture
|
||
def fake_s3_client(monkeypatch):
|
||
_FAKE_S3_CLIENT.clear()
|
||
|
||
def fake_boto3_client(service_name, *args, **kwargs):
|
||
assert service_name == "s3"
|
||
return _FAKE_S3_CLIENT
|
||
|
||
monkeypatch.setattr(boto3, "client", fake_boto3_client)
|
||
litellm.success_callback = []
|
||
litellm.callbacks = []
|
||
yield _FAKE_S3_CLIENT
|
||
litellm.success_callback = []
|
||
litellm.callbacks = []
|
||
|
||
@pytest.mark.usefixtures(
|
||
"cancel_s3_periodic_flush_tasks",
|
||
"fake_s3_client",
|
||
"_vcr_outcome_gate",
|
||
"drain_logging_worker",
|
||
"isolate_litellm_state",
|
||
"setup_and_teardown",
|
||
"amazing_s3_retries",
|
||
)
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.parametrize("sync_mode,streaming", [(True, True), (True, False), (False, True), (False, False)])
|
||
@pytest.mark.flaky(retries=3, delay=1)
|
||
async def test_basic_s3_logging(sync_mode, streaming):
|
||
verbose_logger.setLevel(level=logging.DEBUG)
|
||
litellm.success_callback = ["s3"]
|
||
litellm.s3_callback_params = {
|
||
"s3_bucket_name": "load-testing-oct",
|
||
"s3_aws_secret_access_key": "os.environ/AWS_SECRET_ACCESS_KEY",
|
||
"s3_aws_access_key_id": "os.environ/AWS_ACCESS_KEY_ID",
|
||
"s3_region_name": "us-west-2",
|
||
}
|
||
litellm.set_verbose = True
|
||
response_id = None
|
||
if sync_mode is True:
|
||
response = litellm.completion(
|
||
model="gpt-5-mini",
|
||
messages=[{"role": "user", "content": "This is a test"}],
|
||
mock_response="It's simple to use and easy to get started",
|
||
stream=streaming,
|
||
)
|
||
if streaming:
|
||
for chunk in response:
|
||
response_id = chunk.id
|
||
else:
|
||
response_id = response.id
|
||
assert _FAKE_S3_CLIENT.object_written.wait(timeout=5)
|
||
else:
|
||
response = await litellm.acompletion(
|
||
model="gpt-5-mini",
|
||
messages=[{"role": "user", "content": "This is a test"}],
|
||
mock_response="It's simple to use and easy to get started",
|
||
stream=streaming,
|
||
)
|
||
if streaming:
|
||
async for chunk in response:
|
||
response_id = chunk.id
|
||
else:
|
||
response_id = response.id
|
||
assert await asyncio.to_thread(_FAKE_S3_CLIENT.object_written.wait, 5)
|
||
|
||
total_objects, all_s3_keys = list_all_s3_objects("load-testing-oct")
|
||
|
||
assert any(response_id in key for key in all_s3_keys)
|
||
s3 = boto3.client("s3")
|
||
for key in all_s3_keys:
|
||
s3.delete_object(Bucket="load-testing-oct", Key=key)
|
||
|
||
@pytest.mark.usefixtures(
|
||
"cancel_s3_periodic_flush_tasks",
|
||
"fake_s3_client",
|
||
"_vcr_outcome_gate",
|
||
"drain_logging_worker",
|
||
"isolate_litellm_state",
|
||
"setup_and_teardown",
|
||
"amazing_s3_retries",
|
||
)
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.parametrize("streaming", [True])
|
||
@pytest.mark.flaky(retries=3, delay=1)
|
||
async def test_basic_s3_v2_logging(streaming):
|
||
from litellm.integrations.s3_v2 import S3Logger
|
||
|
||
litellm.s3_callback_params = {
|
||
"s3_bucket_name": "load-testing-oct",
|
||
"s3_aws_secret_access_key": "test-secret",
|
||
"s3_aws_access_key_id": "test-key",
|
||
"s3_region_name": "us-west-2",
|
||
}
|
||
|
||
s3_v2_logger = S3Logger(s3_flush_interval=1)
|
||
litellm.callbacks = [s3_v2_logger]
|
||
|
||
uploaded_keys: list = []
|
||
uploaded = asyncio.Event()
|
||
|
||
async def mock_upload(batch_logging_element):
|
||
uploaded_keys.append(batch_logging_element.s3_object_key)
|
||
uploaded.set()
|
||
|
||
s3_v2_logger.async_upload_data_to_s3 = mock_upload
|
||
|
||
litellm.set_verbose = True
|
||
response_id = None
|
||
response = await litellm.acompletion(
|
||
model="gpt-5-mini",
|
||
messages=[{"role": "user", "content": "This is a test"}],
|
||
mock_response="It's simple to use and easy to get started",
|
||
stream=streaming,
|
||
)
|
||
if streaming:
|
||
async for chunk in response:
|
||
response_id = chunk.id
|
||
else:
|
||
response_id = response.id
|
||
|
||
assert await asyncio.wait_for(uploaded.wait(), timeout=5)
|
||
|
||
assert len(uploaded_keys) > 0, "S3 upload was never called"
|
||
assert any(response_id in key for key in uploaded_keys), (
|
||
f"Expected response_id={response_id} in one of the uploaded S3 keys: {uploaded_keys}"
|
||
)
|
||
|
||
@pytest.mark.usefixtures(
|
||
"cancel_s3_periodic_flush_tasks",
|
||
"fake_s3_client",
|
||
"_vcr_outcome_gate",
|
||
"drain_logging_worker",
|
||
"isolate_litellm_state",
|
||
"setup_and_teardown",
|
||
"amazing_s3_retries",
|
||
)
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.flaky(retries=3, delay=1)
|
||
async def test_basic_s3_v2_logging_failure():
|
||
"""Test that S3 v2 logger makes httpx PUT request when logging failures"""
|
||
from unittest.mock import AsyncMock, MagicMock
|
||
|
||
from litellm.integrations.s3_v2 import S3Logger
|
||
|
||
s3_v2_logger = S3Logger(s3_flush_interval=1)
|
||
|
||
mock_response = MagicMock()
|
||
mock_response.status_code = 200
|
||
mock_response.raise_for_status = MagicMock()
|
||
|
||
s3_v2_logger.async_httpx_client = AsyncMock()
|
||
s3_v2_logger.async_httpx_client.put.return_value = mock_response
|
||
|
||
upload_called = False
|
||
uploaded = asyncio.Event()
|
||
|
||
async def mock_upload(batch_logging_element):
|
||
nonlocal upload_called
|
||
upload_called = True
|
||
url = f"https://test-bucket.s3.us-west-2.amazonaws.com/{batch_logging_element.s3_object_key}"
|
||
headers = {"Content-Type": "application/json"}
|
||
data = '{"model": "gpt-5-mini"}'
|
||
|
||
await s3_v2_logger.async_httpx_client.put(url=url, headers=headers, data=data)
|
||
uploaded.set()
|
||
|
||
s3_v2_logger.async_upload_data_to_s3 = mock_upload
|
||
|
||
litellm.callbacks = [s3_v2_logger]
|
||
litellm.s3_callback_params = {
|
||
"s3_bucket_name": "test-bucket",
|
||
"s3_aws_secret_access_key": "test-secret",
|
||
"s3_aws_access_key_id": "test-key",
|
||
"s3_region_name": "us-west-2",
|
||
}
|
||
litellm.set_verbose = True
|
||
|
||
try:
|
||
await litellm.acompletion(
|
||
model="gpt-5-mini",
|
||
api_key="invalid-api-key",
|
||
messages=[{"role": "user", "content": "This is a test"}],
|
||
mock_response=Exception("forced failure for S3 logging test"),
|
||
)
|
||
except Exception:
|
||
pass
|
||
|
||
assert await asyncio.wait_for(uploaded.wait(), timeout=5)
|
||
|
||
assert upload_called, "S3 upload method was not called"
|
||
|
||
s3_v2_logger.async_httpx_client.put.assert_called()
|
||
|
||
call_args = s3_v2_logger.async_httpx_client.put.call_args
|
||
assert call_args is not None
|
||
url = call_args[1]["url"] if "url" in call_args[1] else call_args[0][0]
|
||
|
||
assert "test-bucket.s3.us-west-2.amazonaws.com" in url
|
||
|
||
headers = call_args[1]["headers"]
|
||
assert headers["Content-Type"] == "application/json"
|
||
|
||
data = call_args[1]["data"]
|
||
assert data is not None
|
||
assert '"model": "gpt-5-mini"' in data
|
||
|
||
def list_all_s3_objects(bucket_name):
|
||
s3 = boto3.client("s3")
|
||
|
||
all_s3_keys = []
|
||
|
||
paginator = s3.get_paginator("list_objects_v2")
|
||
total_objects = 0
|
||
|
||
for page in paginator.paginate(Bucket=bucket_name):
|
||
if "Contents" in page:
|
||
total_objects += len(page["Contents"])
|
||
all_s3_keys.extend([obj["Key"] for obj in page["Contents"]])
|
||
|
||
return total_objects, all_s3_keys
|
||
|
||
class TestS3Logger(S3Logger):
|
||
def __init__(self, *args, **kwargs):
|
||
self.recorded_requests = {}
|
||
self.logged_standard_logging_payload = None
|
||
super().__init__(*args, **kwargs)
|
||
|
||
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||
self.recorded_requests[response_obj["id"]] = start_time
|
||
self.logged_standard_logging_payload = kwargs["standard_logging_object"]
|
||
return await super().async_log_success_event(kwargs, response_obj, start_time, end_time)
|