mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-25 01:02:15 +00:00
test(s3_v2): add coverage for sync upload and download URL encoding
Add tests verifying URL-encoding of special characters in both sync_upload_data_to_s3 and _download_object_from_s3 paths.
This commit is contained in:
parent
41e054c798
commit
b5b663731e
1 changed files with 88 additions and 0 deletions
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue