From d63ceb9eb3b1e1660807a2df02c5329e6a8887dd Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 7 Mar 2026 16:44:35 -0800 Subject: [PATCH] fix(tests): fix image variation fixture and SSE parsing in streaming cost injection tests - image_variation: replace non-square S3 URL fetch with programmatic 1024x1024 RGBA PNG using Pillow. DALL-E 2 requires a square PNG for create_variation. - anthropic passthrough streaming tests: split each HTTP chunk by newlines before checking for 'data: ' prefix. The AnthropicResponsesStreamWrapper (OpenAI models path) yields full multi-line SSE events as single bytes objects, so the old line-by-line check missed them. --- tests/image_gen_tests/test_image_variation.py | 19 ++--- .../test_anthropic_passthrough.py | 80 +++++++++++-------- 2 files changed, 53 insertions(+), 46 deletions(-) diff --git a/tests/image_gen_tests/test_image_variation.py b/tests/image_gen_tests/test_image_variation.py index d4f66603352..c0ad5d38e66 100644 --- a/tests/image_gen_tests/test_image_variation.py +++ b/tests/image_gen_tests/test_image_variation.py @@ -27,24 +27,21 @@ import tempfile from base_image_generation_test import BaseImageGenTest import logging from litellm._logging import verbose_logger -import requests from io import BytesIO +from PIL import Image as PILImage verbose_logger.setLevel(logging.DEBUG) @pytest.fixture def image_url(): - # URL of the image - image_url = "https://litellm-listing.s3.amazonaws.com/litellm_logo.png" - - # Fetch the image from the URL - response = requests.get(image_url) - print(response) - response.raise_for_status() # Ensure the request was successful - - # Load the image into a file-like object - image_file = BytesIO(response.content) + # DALL-E 2 image variations require a square PNG (less than 4MB) + # Generate a 1024x1024 square PNG programmatically to avoid network dependency + # and the non-square aspect ratio of the old LiteLLM logo URL + img = PILImage.new("RGBA", (1024, 1024), color=(128, 128, 128, 255)) + image_file = BytesIO() + img.save(image_file, format="PNG") + image_file.seek(0) return image_file diff --git a/tests/pass_through_tests/test_anthropic_passthrough.py b/tests/pass_through_tests/test_anthropic_passthrough.py index e229c08f6e4..253ce70d8e3 100644 --- a/tests/pass_through_tests/test_anthropic_passthrough.py +++ b/tests/pass_through_tests/test_anthropic_passthrough.py @@ -342,35 +342,40 @@ async def test_anthropic_messages_streaming_cost_injection(): headers=headers ) as response: assert response.status == 200 - + # Collect all SSE events + # Split each chunk by newlines to handle both: + # - Anthropic direct path: chunks arrive as individual lines + # - OpenAI/Responses API path: chunks are full multi-line SSE events events = [] - async for line in response.content: - line_str = line.decode('utf-8').strip() - if line_str.startswith('data: '): - try: - data = json.loads(line_str[6:]) # Remove 'data: ' prefix - events.append(data) - except json.JSONDecodeError: - continue - + async for chunk in response.content: + chunk_str = chunk.decode('utf-8') + for line in chunk_str.split('\n'): + line = line.strip() + if line.startswith('data: '): + try: + data = json.loads(line[6:]) # Remove 'data: ' prefix + events.append(data) + except json.JSONDecodeError: + continue + # Find message_delta event with usage message_delta_events = [ - event for event in events + event for event in events if event.get('type') == 'message_delta' and 'usage' in event ] - + assert len(message_delta_events) > 0, "No message_delta events with usage found" - + # Check that cost is included in usage for event in message_delta_events: usage = event.get('usage', {}) assert 'cost' in usage, f"Cost not found in usage: {usage}" assert isinstance(usage['cost'], (int, float)), f"Cost should be numeric: {usage['cost']}" assert usage['cost'] >= 0, f"Cost should be non-negative: {usage['cost']}" - + print(f"✅ Found message_delta with cost: {usage}") - + print(f"✅ Test passed: Found {len(message_delta_events)} message_delta events with cost") @@ -381,54 +386,59 @@ async def test_anthropic_messages_openai_model_streaming_cost_injection(): Test that cost is injected into message_delta usage for OpenAI model via Anthropic Messages API """ print("Testing cost injection in Anthropic Messages API with OpenAI model") - + headers = { "Authorization": "Bearer sk-1234", "Content-Type": "application/json", "anthropic-version": "2023-06-01", } - + payload = { "model": "openai/gpt-4o", "max_tokens": 10, "stream": True, "messages": [{"role": "user", "content": "Say 'Hi'"}], } - + async with aiohttp.ClientSession() as session: async with session.post( - "http://0.0.0.0:4000/v1/messages", - json=payload, + "http://0.0.0.0:4000/v1/messages", + json=payload, headers=headers ) as response: assert response.status == 200 - + # Collect all SSE events + # Split each chunk by newlines to handle both: + # - Direct API paths: chunks arrive as individual lines + # - OpenAI/Responses API path: chunks are full multi-line SSE events events = [] - async for line in response.content: - line_str = line.decode('utf-8').strip() - if line_str.startswith('data: '): - try: - data = json.loads(line_str[6:]) # Remove 'data: ' prefix - events.append(data) - except json.JSONDecodeError: - continue - + async for chunk in response.content: + chunk_str = chunk.decode('utf-8') + for line in chunk_str.split('\n'): + line = line.strip() + if line.startswith('data: '): + try: + data = json.loads(line[6:]) # Remove 'data: ' prefix + events.append(data) + except json.JSONDecodeError: + continue + # Find message_delta event with usage message_delta_events = [ - event for event in events + event for event in events if event.get('type') == 'message_delta' and 'usage' in event ] - + assert len(message_delta_events) > 0, "No message_delta events with usage found" - + # Check that cost is included in usage for event in message_delta_events: usage = event.get('usage', {}) assert 'cost' in usage, f"Cost not found in usage: {usage}" assert isinstance(usage['cost'], (int, float)), f"Cost should be numeric: {usage['cost']}" assert usage['cost'] >= 0, f"Cost should be non-negative: {usage['cost']}" - + print(f"✅ Found message_delta with cost: {usage}") - + print(f"✅ Test passed: Found {len(message_delta_events)} message_delta events with cost")