litellm/tests/test_litellm/integrations/test_s3.py

from datetime import datetime
from unittest.mock import MagicMock, patch

import litellm
from litellm.constants import MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES, MAX_S3_OBJECT_KEY_BYTES
from litellm.integrations.s3 import S3Logger

TEST_KMS_KEY_ARN = "arn:aws:kms:us-east-1:111122223333:key/test-key-id"


def _standard_logging_payload(response_id: str = "chatcmpl-test-id") -> dict:
    return {
        "id": response_id,
        "metadata": {"user_api_key_team_alias": None},
    }


def _log_event_kwargs(response_id: str = "chatcmpl-test-id") -> dict:
    return {
        "litellm_params": {"metadata": {}},
        "standard_logging_object": _standard_logging_payload(response_id),
    }


def _run_log_event(callback_params: dict, response_id: str = "chatcmpl-test-id") -> MagicMock:
    original = litellm.s3_callback_params
    litellm.s3_callback_params = callback_params
    try:
        with patch("boto3.client") as mock_boto3_client:
            mock_s3_client = MagicMock()
            mock_boto3_client.return_value = mock_s3_client
            logger = S3Logger()
            logger.log_event(
                kwargs=_log_event_kwargs(response_id),
                response_obj={"id": response_id},
                start_time=datetime(2026, 7, 30, 12, 0, 0),
                end_time=datetime(2026, 7, 30, 12, 0, 1),
                print_verbose=lambda *args, **kwargs: None,
            )
        return mock_s3_client
    finally:
        litellm.s3_callback_params = original


def test_put_object_includes_sse_kms_params_when_configured():
    """
    When s3_server_side_encryption and s3_sse_kms_key_id are set in
    s3_callback_params, put_object must receive ServerSideEncryption and
    SSEKMSKeyId so objects land encrypted with the customer-managed key.
    """
    mock_s3_client = _run_log_event(
        {
            "s3_bucket_name": "test-bucket",
            "s3_region_name": "us-east-1",
            "s3_server_side_encryption": "aws:kms",
            "s3_sse_kms_key_id": TEST_KMS_KEY_ARN,
        }
    )

    put_object_kwargs = mock_s3_client.put_object.call_args.kwargs
    assert put_object_kwargs["ServerSideEncryption"] == "aws:kms"
    assert put_object_kwargs["SSEKMSKeyId"] == TEST_KMS_KEY_ARN


def test_put_object_supports_sse_s3_without_key_id():
    """SSE-S3 (AES256) needs only ServerSideEncryption, no key id."""
    mock_s3_client = _run_log_event(
        {
            "s3_bucket_name": "test-bucket",
            "s3_region_name": "us-east-1",
            "s3_server_side_encryption": "AES256",
        }
    )

    put_object_kwargs = mock_s3_client.put_object.call_args.kwargs
    assert put_object_kwargs["ServerSideEncryption"] == "AES256"
    assert "SSEKMSKeyId" not in put_object_kwargs


def test_put_object_omits_sse_params_by_default():
    """Without SSE config, put_object kwargs must stay unchanged."""
    mock_s3_client = _run_log_event(
        {
            "s3_bucket_name": "test-bucket",
            "s3_region_name": "us-east-1",
        }
    )

    put_object_kwargs = mock_s3_client.put_object.call_args.kwargs
    assert "ServerSideEncryption" not in put_object_kwargs
    assert "SSEKMSKeyId" not in put_object_kwargs


def test_put_object_infers_aws_kms_when_only_key_id_set():
    """A key id without an algorithm must infer aws:kms instead of sending an invalid request."""
    mock_s3_client = _run_log_event(
        {
            "s3_bucket_name": "test-bucket",
            "s3_region_name": "us-east-1",
            "s3_sse_kms_key_id": TEST_KMS_KEY_ARN,
        }
    )

    put_object_kwargs = mock_s3_client.put_object.call_args.kwargs
    assert put_object_kwargs["ServerSideEncryption"] == "aws:kms"
    assert put_object_kwargs["SSEKMSKeyId"] == TEST_KMS_KEY_ARN


def test_put_object_drops_key_id_when_algorithm_is_not_kms():
    """AES256 plus a key id is invalid for S3; the key id must be dropped, not sent."""
    mock_s3_client = _run_log_event(
        {
            "s3_bucket_name": "test-bucket",
            "s3_region_name": "us-east-1",
            "s3_server_side_encryption": "AES256",
            "s3_sse_kms_key_id": TEST_KMS_KEY_ARN,
        }
    )

    put_object_kwargs = mock_s3_client.put_object.call_args.kwargs
    assert put_object_kwargs["ServerSideEncryption"] == "AES256"
    assert "SSEKMSKeyId" not in put_object_kwargs


def test_non_string_algorithm_is_dropped_and_valid_key_id_is_rescued():
    """
    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.
    """
    mock_s3_client = _run_log_event(
        {
            "s3_bucket_name": "test-bucket",
            "s3_region_name": "us-east-1",
            "s3_server_side_encryption": True,
            "s3_sse_kms_key_id": TEST_KMS_KEY_ARN,
        }
    )

    put_object_kwargs = mock_s3_client.put_object.call_args.kwargs
    assert put_object_kwargs["ServerSideEncryption"] == "aws:kms"
    assert put_object_kwargs["SSEKMSKeyId"] == TEST_KMS_KEY_ARN


def test_non_string_key_id_is_dropped_and_valid_algorithm_is_kept():
    """A mistyped key id (unquoted YAML number) must not disable the valid algorithm."""
    mock_s3_client = _run_log_event(
        {
            "s3_bucket_name": "test-bucket",
            "s3_region_name": "us-east-1",
            "s3_server_side_encryption": "aws:kms",
            "s3_sse_kms_key_id": 12345,
        }
    )

    put_object_kwargs = mock_s3_client.put_object.call_args.kwargs
    assert put_object_kwargs["ServerSideEncryption"] == "aws:kms"
    assert "SSEKMSKeyId" not in put_object_kwargs


def test_put_object_key_and_filename_are_bounded_for_an_oversized_response_id():
    """The sync logger bounds both the key and the Content-Disposition filename."""
    mock_s3_client = _run_log_event(
        {"s3_bucket_name": "test-bucket", "s3_region_name": "us-west-2", "s3_path": "logs"},
        response_id="resp_" + "A" * 1100,
    )

    put_object_kwargs = mock_s3_client.put_object.call_args.kwargs
    assert len(put_object_kwargs["Key"].encode("utf-8")) <= MAX_S3_OBJECT_KEY_BYTES
    assert put_object_kwargs["Key"].startswith("logs/2026-07-30/time-12-00-00-000000_resp_")
    filename = put_object_kwargs["ContentDisposition"].removeprefix('inline; filename="').removesuffix('"')
    assert len(filename.encode("utf-8")) <= MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES


def test_put_object_keeps_the_configured_path_intact_when_only_the_id_has_to_shrink():
    """A long configured s3_path survives whole when the id can be shortened instead."""
    long_path = "litellm-prod-logs/" + "t" * 921
    mock_s3_client = _run_log_event(
        {"s3_bucket_name": "test-bucket", "s3_region_name": "us-west-2", "s3_path": long_path},
        response_id="resp_" + "B" * 100,
    )

    key = mock_s3_client.put_object.call_args.kwargs["Key"]
    assert key.startswith(long_path + "/2026-07-30/")
    assert len(key.encode("utf-8")) == MAX_S3_OBJECT_KEY_BYTES