litellm/tests/unit/integrations/test_s3_v2.py
devin-ai-integration[bot] e47b1f2a3f
fix(s3_v2): upload fresh events first, drop terminal failures and hour-old retries by default, opt-in adaptive concurrency (#43022)
* fix(s3_v2): drop terminal upload failures, bound retries per flush and enforce the queue cap at enqueue

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(s3_v2): keep retrying credential-rotation 403s, only AccessDenied-style errors are terminal

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(env_keys): exclude DEFAULT_S3_MAX_FLUSH_ATTEMPTS as an internal tuning var

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(s3_v2): retry every 5xx, warn on first queue overflow, validate the flush budget

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(s3_v2): read the queue cap defensively so un-initialized loggers still enqueue

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(s3_v2): drop the getattr in _enqueue and tighten the retry tests

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(s3_v2): keep the constructor flush budget when the callback override is invalid

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* feat(s3_v2): adapt per-object upload concurrency to sink latency and throttling

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* refactor(s3_v2): tidy adaptive limiter

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* perf(s3_v2): wake one waiter per released upload slot

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* feat(s3_v2): make the enqueue queue cap configurable with s3_max_queue_size

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(s3_v2): audit cells for cache hits, coded 403 and callback modes

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(s3_v2): retry bucket-wide failures by default, age-budget requeues and make terminal drops and adaptive concurrency opt-in

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* refactor(s3_v2): suppress the missing-waiter ValueError explicitly in the adaptive limiter

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(s3_v2): count oldest events trimmed after a failed flush as callback failures

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* style(s3_v2): move the mutable-ok marker onto the list literal it suppresses

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(s3_v2): report post-flush overflow drops once and grow adaptive concurrency above the floor before asserting back-off

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(s3_v2): fail the SlowDown back-off test when the measured window sees no PUTs

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(s3_v2): restore base retry defaults, opt-in age budget, no enqueue cap, back off outside the limiter slot

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(s3_v2): hoist the default no-op upload slot to a module constant

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(s3_v2): drop unused mutable-ok suppressions on queue appends

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(s3_v2): keep the retry queue oldest-first and prioritise fresh events at upload time

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* style(s3_v2): keep the mutable-ok marker on the queue list literal

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(s3_v2): build request bodies inside the upload slot and keep the sync retry set at base parity

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(s3_v2): drop wall-clock sleeps from the unit tests

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(s3_v2): rebuild the request body inside the slot on every retry attempt

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(s3_v2): anchor the backoff window on the first observed failure and tighten shard assertions

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(s3_v2): default the upload slot to the logger limiter so monkeypatched doubles keep working

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(s3_v2): clear ambient AWS env credentials so the rotating profile signs the sync retry test

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(s3_v2): shrink the linear send-batch perf test to 2k/8k elements

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(s3_v2): fix stale batch sizes in the perf test assert message

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(s3_v2): keep async in-call retries on the base 403/500/503 set

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(s3_v2): hold the upload slot across retries, restore the bool upload contract, and fail safe on bool config

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(s3_v2): mark dropped uploads by element identity so a shared key cannot mask a retryable sibling

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(s3_v2): make the per-flush drop lookup constant time

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(s3_v2): take the upload slot in the caller like base, build the body once per attempt loop

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(s3_v2): match base retry, logging and hook behaviour unless the new options are opted in

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(s3_v2): assert the wire key in the init-bypassed sync upload test

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(s3_v2): assert the signed headers and wire key in the init-bypassed sync upload test

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(s3_v2): default the upload limiter at class level instead of reading it with getattr

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(s3_v2): drop the duplicate annotations that redeclare the class-level counters

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* feat(s3_v2): drop terminal-failed uploads by default and bound retry age to one hour

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(s3_v2): fall back to the configured retry age on invalid values

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(s3_v2): drive retry-age tests from a fixed clock

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(s3_v2): audit cells for retry-age opt-out and 429 single-put parity

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: yucheng <yucheng@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
2026-09-26 14:58:28 -07:00

4794 lines
168 KiB
Python
Raw Blame History

This file contains invisible Unicode characters

This file contains invisible Unicode characters that are indistinguishable to humans but may be processed differently by a computer. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import asyncio
import copy
import json
import re
import sys
import textwrap
import time
import uuid
from collections.abc import Awaitable, Callable
from contextlib import asynccontextmanager
from datetime import datetime
from pathlib import Path
from typing import Final
from unittest.mock import AsyncMock, MagicMock, call, patch
import httpx
import pytest
import respx
from litellm.integrations.s3_v2 import S3BatchUploadError, S3Logger
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.types.integrations.s3_v2 import s3BatchLoggingElement
from litellm.types.utils import StandardLoggingPayload
_real_sleep: Final = asyncio.sleep
_NOW: Final = 1_000_000.0
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
# Verify sleep was called with the backoff delay
mock_sleep.assert_called_once_with(1) # 2**0 = 1s
@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
mock_sleep.assert_called_once_with(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]):
"""
S3Logger wired to the real AsyncHTTPHandler over an httpx MockTransport that answers with the given
statuses in order, so the handler's own raise_for_status behaviour is exercised end to end.
"""
requests: list[httpx.Request] = []
replies = iter(statuses)
def respond(request: httpx.Request) -> httpx.Response:
requests.append(request)
return httpx.Response(next(replies), 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"]
mock_sleep.assert_awaited_once_with(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 mock_sleep.await_args_list == [call(1), call(2)]
assert "Error uploading to s3" 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 mock_sleep.await_args_list == [call(1), call(2)]
assert "Error uploading to s3" in caplog.text
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)
mock_sleep.assert_called_once_with(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
mock_sleep.assert_called_once_with(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.fixture(autouse=True)
def patch_asyncio_create_task():
"""Prevent 'no running event loop' errors when S3Logger calls asyncio.create_task()."""
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)
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) -> 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) -> 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) -> 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) -> 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) -> 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) -> 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) -> 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) -> 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_sleeps"),
[
pytest.param(429, 1, [], id="429-single"),
pytest.param(408, 1, [], id="408-single"),
pytest.param(502, 1, [], id="502-single"),
pytest.param(504, 1, [], id="504-single"),
pytest.param(503, 3, [call(1), call(2)], id="503-backoff"),
],
)
def test_sync_upload_retry_set_matches_base(status: int, expected_puts: int, expected_sleeps: list) -> 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 mock_sleep.call_args_list == expected_sleeps
@pytest.mark.asyncio
@pytest.mark.parametrize(
("status", "code"),
[
pytest.param(503, "SlowDown", id="503"),
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 mock_sleep.await_args_list == [call(1), call(2)]
assert len(logger.log_queue) == 1
@pytest.mark.asyncio
@pytest.mark.parametrize(
("status", "code"),
[
pytest.param(429, "TooManyRequests", id="429"),
pytest.param(408, None, id="408"),
pytest.param(502, None, id="502"),
pytest.param(504, None, id="504"),
],
)
async def test_non_base_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) -> 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) -> 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) -> 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) -> 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) -> 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 mock_sleep.call_args_list == [call(1), call(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) -> 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) -> 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) -> 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"