diff --git a/tests/test_litellm/integrations/test_s3_v2.py b/tests/test_litellm/integrations/test_s3_v2.py index 2cfa0b0612b..ccf9b8d2ef6 100644 --- a/tests/test_litellm/integrations/test_s3_v2.py +++ b/tests/test_litellm/integrations/test_s3_v2.py @@ -1242,3 +1242,91 @@ async def test_s3_v2_put_url_encodes_special_chars( assert " " not in actual_url assert "#" not in actual_url assert "/" in actual_url + + +@pytest.mark.asyncio +@patch("asyncio.create_task") +@patch("litellm.integrations.s3_v2.CustomBatchLogger.periodic_flush") +async def test_s3_v2_download_url_encodes_special_chars( + mock_periodic_flush, mock_create_task +): + """GET path: URL-encode special characters in s3_object_key for download.""" + from urllib.parse import quote + from unittest.mock import AsyncMock + + mock_periodic_flush.return_value = None + mock_create_task.return_value = None + + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = {"test": "data"} + + s3_object_key = "team α/logs/2024-01-01 12:00#special.json" + + s3_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_logger.async_httpx_client = AsyncMock() + s3_logger.async_httpx_client.get.return_value = mock_response + + await s3_logger._download_object_from_s3(s3_object_key) + + call_args = s3_logger.async_httpx_client.get.call_args + assert call_args is not None + actual_url = call_args[0][0] + encoded_key = quote(s3_object_key, safe="/") + expected_url = f"https://test-bucket.s3.us-east-1.amazonaws.com/{encoded_key}" + assert actual_url == expected_url + assert " " not in actual_url + assert "#" not in actual_url + + +@patch("asyncio.create_task") +@patch("litellm.integrations.s3_v2.CustomBatchLogger.periodic_flush") +def test_s3_v2_sync_upload_url_encodes_special_chars( + mock_periodic_flush, mock_create_task +): + """Sync PUT path: URL-encode special characters in s3_object_key.""" + from urllib.parse import quote + from unittest.mock import MagicMock as MM + + 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 = "team α/logs/2024-01-01 12:00#special.json" + test_element = s3BatchLoggingElement( + s3_object_key=s3_object_key, + payload={"test": "data"}, + s3_object_download_filename="special.json", + ) + + s3_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_httpx = MagicMock() + mock_httpx.put.return_value = mock_response + + with patch("litellm.integrations.s3_v2._get_httpx_client", return_value=mock_httpx): + s3_logger.sync_upload_data_to_s3(test_element) + + call_args = mock_httpx.put.call_args + assert call_args is not None + actual_url = call_args[0][0] + encoded_key = quote(s3_object_key, safe="/") + expected_url = f"https://test-bucket.s3.us-east-1.amazonaws.com/{encoded_key}" + assert actual_url == expected_url + assert " " not in actual_url + assert "#" not in actual_url