mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-11 17:05:43 +00:00
Ishaan Jaffer
95 lines
2.9 KiB
Python
95 lines
2.9 KiB
Python
import asyncio
|
|
import os
|
|
import sys
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
sys.path.insert(0, os.path.abspath("../.."))
|
|
|
|
import litellm
|
|
from litellm import aimage_generation
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"model,expected_endpoint",
|
|
[
|
|
("fal_ai/fal-ai/flux-pro/v1.1-ultra", "fal-ai/flux-pro/v1.1-ultra"),
|
|
("fal_ai/fal-ai/stable-diffusion-v35-medium", "fal-ai/stable-diffusion-v35-medium"),
|
|
],
|
|
)
|
|
@pytest.mark.asyncio
|
|
async def test_fal_ai_image_generation_basic(model, expected_endpoint):
|
|
"""
|
|
Test that fal_ai image generation constructs correct request body and URL.
|
|
|
|
Validates:
|
|
- Correct API endpoint URL construction
|
|
- Proper request body format with prompt
|
|
- Correct Authorization header format
|
|
"""
|
|
captured_url = None
|
|
captured_json_data = None
|
|
captured_headers = None
|
|
|
|
def capture_post_call(*args, **kwargs):
|
|
nonlocal captured_url, captured_json_data, captured_headers
|
|
|
|
captured_url = args[0] if args else kwargs.get("url")
|
|
captured_json_data = kwargs.get("json")
|
|
captured_headers = kwargs.get("headers")
|
|
|
|
# Mock response with fal.ai format
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 200
|
|
mock_response.headers = {"content-type": "application/json"}
|
|
mock_response.json.return_value = {
|
|
"images": [
|
|
{
|
|
"url": "https://example.com/generated-image.png",
|
|
"width": 1024,
|
|
"height": 768,
|
|
"content_type": "image/jpeg"
|
|
}
|
|
],
|
|
"seed": 42
|
|
}
|
|
|
|
return mock_response
|
|
|
|
with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") as mock_post:
|
|
mock_post.side_effect = capture_post_call
|
|
|
|
test_api_key = "test-fal-ai-key-12345"
|
|
test_prompt = "A cute baby sea otter"
|
|
|
|
response = await aimage_generation(
|
|
model=model,
|
|
prompt=test_prompt,
|
|
api_key=test_api_key,
|
|
)
|
|
|
|
# Validate response
|
|
assert response is not None
|
|
assert hasattr(response, "data")
|
|
assert response.data is not None
|
|
assert len(response.data) > 0
|
|
|
|
# Validate URL
|
|
assert captured_url is not None
|
|
assert "fal.run" in captured_url
|
|
assert expected_endpoint in captured_url
|
|
print(f"Validated URL: {captured_url}")
|
|
|
|
# Validate headers
|
|
assert captured_headers is not None
|
|
assert "Authorization" in captured_headers
|
|
assert captured_headers["Authorization"] == f"Key {test_api_key}"
|
|
print(f"Validated headers: {captured_headers}")
|
|
|
|
# Validate request body
|
|
assert captured_json_data is not None
|
|
assert captured_json_data["prompt"] == test_prompt
|
|
print(f"Validated request body: {captured_json_data}")
|
|
|