scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 124

fix(tests): fix image variation fixture and SSE parsing in streaming cost injection tests

Browse source 
- 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.
This commit is contained in:
Ishaan Jaffer 2026-03-07 16:44:35 -08:00
parent a50a84c16c
commit d63ceb9eb3
2 changed files with 53 additions and 46 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

19
tests/image_gen_tests/test_image_variation.py
View file

@ -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

80
tests/pass_through_tests/test_anthropic_passthrough.py
View file

@ -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")