diff --git a/tests/litellm/llms/sagemaker/test_sagemaker_common_utils.py b/tests/litellm/llms/sagemaker/test_sagemaker_common_utils.py new file mode 100644 index 00000000000..e19c41f1ac3 --- /dev/null +++ b/tests/litellm/llms/sagemaker/test_sagemaker_common_utils.py @@ -0,0 +1,93 @@ +import json +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest + +sys.path.insert(0, os.path.abspath("../../../../..")) +from litellm.llms.sagemaker.common_utils import AWSEventStreamDecoder + + +@pytest.mark.asyncio +async def test_aiter_bytes_unicode_decode_error(): + """ + Test that AWSEventStreamDecoder.aiter_bytes() does not raise an error when encountering invalid UTF-8 bytes. (UnicodeDecodeError) + + + Ensures stream processing continues despite the error. + """ + # Create an instance of AWSEventStreamDecoder + decoder = AWSEventStreamDecoder(model="test-model") + + # Create a mock event that will trigger a UnicodeDecodeError + mock_event = MagicMock() + mock_event.to_response_dict.return_value = { + "status_code": 200, + "headers": {}, + "body": b"\xff\xfe", # Invalid UTF-8 bytes + } + + # Create a mock EventStreamBuffer that yields our mock event + mock_buffer = MagicMock() + mock_buffer.__iter__.return_value = [mock_event] + + # Mock the EventStreamBuffer class + with patch("botocore.eventstream.EventStreamBuffer", return_value=mock_buffer): + # Create an async generator that yields some test bytes + async def mock_iterator(): + yield b"" + + # Process the stream + chunks = [] + async for chunk in decoder.aiter_bytes(mock_iterator()): + if chunk is not None: + print("chunk=", chunk) + chunks.append(chunk) + + # Verify that processing continued despite the error + # The chunks list should be empty since we only sent invalid data + assert len(chunks) == 0 + + +@pytest.mark.asyncio +async def test_aiter_bytes_valid_chunk_followed_by_unicode_error(): + """ + Test that valid chunks are processed correctly even when followed by Unicode decode errors. + This ensures errors don't corrupt or prevent processing of valid data that came before. + """ + decoder = AWSEventStreamDecoder(model="test-model") + + # Create two mock events - first valid, then invalid + mock_valid_event = MagicMock() + mock_valid_event.to_response_dict.return_value = { + "status_code": 200, + "headers": {}, + "body": json.dumps({"token": {"text": "hello"}}).encode(), # Valid data first + } + + mock_invalid_event = MagicMock() + mock_invalid_event.to_response_dict.return_value = { + "status_code": 200, + "headers": {}, + "body": b"\xff\xfe", # Invalid UTF-8 bytes second + } + + # Create a mock EventStreamBuffer that yields valid event first, then invalid + mock_buffer = MagicMock() + mock_buffer.__iter__.return_value = [mock_valid_event, mock_invalid_event] + + with patch("botocore.eventstream.EventStreamBuffer", return_value=mock_buffer): + + async def mock_iterator(): + yield b"test_bytes" + + chunks = [] + async for chunk in decoder.aiter_bytes(mock_iterator()): + if chunk is not None: + chunks.append(chunk) + + # Verify we got our valid chunk despite the subsequent error + assert len(chunks) == 1 + assert chunks[0]["text"] == "hello" # Verify the content of the valid chunk