mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-11 01:05:19 +00:00
Ishaan Jaff
fdaad51015
* feat: add add azure endpoint for image endpoints * test: azure image routes working as expected * test azure routes
120 lines
3.7 KiB
Python
120 lines
3.7 KiB
Python
import asyncio
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
from unittest import mock
|
|
|
|
import pytest
|
|
from fastapi.testclient import TestClient
|
|
|
|
sys.path.insert(0, os.path.abspath("../../.."))
|
|
import litellm
|
|
from litellm.proxy.proxy_server import app, initialize
|
|
|
|
example_image_generation_result = {
|
|
"created": 1589478378,
|
|
"data": [{"url": "https://example.com/image.png"}],
|
|
}
|
|
|
|
example_image_edit_result = {
|
|
"created": 1589478400,
|
|
"data": [
|
|
{
|
|
"b64_json": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg=="
|
|
}
|
|
],
|
|
}
|
|
|
|
|
|
def mock_patch_aimage_generation():
|
|
"""Patch the underlying image generation call used by the Router."""
|
|
mock_obj = mock.AsyncMock(return_value=example_image_generation_result)
|
|
mock_obj.__name__ = "aimage_generation"
|
|
return mock.patch(
|
|
"litellm.aimage_generation",
|
|
new_callable=lambda: mock_obj,
|
|
)
|
|
|
|
|
|
def mock_patch_aimage_edit():
|
|
"""Patch the underlying image edit call used by the Router."""
|
|
mock_obj = mock.AsyncMock(return_value=example_image_edit_result)
|
|
mock_obj.__name__ = "aimage_edit"
|
|
return mock.patch(
|
|
"litellm.aimage_edit",
|
|
new_callable=lambda: mock_obj,
|
|
)
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def client_no_auth():
|
|
from litellm.proxy.proxy_server import cleanup_router_config_variables
|
|
|
|
cleanup_router_config_variables()
|
|
repo_root = Path(__file__).resolve().parents[4]
|
|
config_fp = (
|
|
repo_root
|
|
/ "tests"
|
|
/ "proxy_unit_tests"
|
|
/ "test_configs"
|
|
/ "test_config_no_auth.yaml"
|
|
)
|
|
config_fp = str(config_fp)
|
|
|
|
# Create mock objects with __name__ attribute
|
|
mock_generation = mock.AsyncMock(return_value=example_image_generation_result)
|
|
mock_generation.__name__ = "aimage_generation"
|
|
|
|
mock_edit = mock.AsyncMock(return_value=example_image_edit_result)
|
|
mock_edit.__name__ = "aimage_edit"
|
|
|
|
with mock.patch(
|
|
"litellm.aimage_generation",
|
|
new_callable=lambda: mock_generation,
|
|
) as patched_generation, mock.patch(
|
|
"litellm.aimage_edit",
|
|
new_callable=lambda: mock_edit,
|
|
) as patched_edit:
|
|
asyncio.run(initialize(config=config_fp, debug=True))
|
|
client = TestClient(app)
|
|
yield client, patched_generation, patched_edit
|
|
|
|
|
|
def test_azure_image_generation_route(client_no_auth):
|
|
client, mock_aimage_generation, _ = client_no_auth
|
|
test_data = {"prompt": "A cute baby sea otter", "n": 1, "size": "1024x1024"}
|
|
response = client.post(
|
|
"/openai/deployments/dall-e-3/images/generations", json=test_data
|
|
)
|
|
|
|
mock_aimage_generation.assert_called_once()
|
|
call_kwargs = mock_aimage_generation.call_args.kwargs
|
|
assert "dall-e-3" in call_kwargs["model"]
|
|
assert call_kwargs["prompt"] == "A cute baby sea otter"
|
|
assert call_kwargs["n"] == 1
|
|
assert call_kwargs["size"] == "1024x1024"
|
|
assert response.status_code == 200
|
|
assert response.json()["data"]
|
|
|
|
|
|
def test_azure_image_edit_route(client_no_auth):
|
|
litellm._turn_on_debug()
|
|
client, _, mock_aimage_edit = client_no_auth
|
|
image_path = os.path.join(
|
|
os.path.dirname(__file__),
|
|
"../../../image_gen_tests/test_image.png",
|
|
)
|
|
with open(image_path, "rb") as f:
|
|
files = {"image": ("test_image.png", f, "image/png")}
|
|
data = {"prompt": "A cute baby sea otter"}
|
|
response = client.post(
|
|
"/openai/deployments/dall-e-3/images/edits", files=files, data=data
|
|
)
|
|
|
|
mock_aimage_edit.assert_called_once()
|
|
called_kwargs = mock_aimage_edit.call_args.kwargs
|
|
assert "dall-e-3" in called_kwargs["model"]
|
|
assert called_kwargs["prompt"] == "A cute baby sea otter"
|
|
assert response.status_code == 200
|
|
assert response.json()["data"]
|