Added LiteLLM to the stack
This commit is contained in:
@@ -0,0 +1,177 @@
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from io import BufferedReader, BytesIO
|
||||
from typing import Dict, List
|
||||
from unittest.mock import MagicMock, mock_open, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../../../../..")
|
||||
) # Adds the parent directory to the system path
|
||||
|
||||
from litellm.llms.recraft.image_edit.transformation import RecraftImageEditConfig
|
||||
from litellm.types.images.main import ImageEditOptionalRequestParams
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.utils import ImageObject, ImageResponse
|
||||
|
||||
|
||||
class TestRecraftImageEditTransformation:
|
||||
"""
|
||||
Unit tests for Recraft image edit transformation functionality.
|
||||
"""
|
||||
|
||||
def setup_method(self):
|
||||
"""Set up test fixtures before each test method."""
|
||||
self.config = RecraftImageEditConfig()
|
||||
self.model = "recraft-v3"
|
||||
self.logging_obj = MagicMock()
|
||||
self.prompt = "Add more trees to this landscape"
|
||||
|
||||
def test_transform_image_edit_request(self):
|
||||
"""
|
||||
Test that transform_image_edit_request correctly transforms request parameters
|
||||
and separates files from data.
|
||||
"""
|
||||
# Mock image data
|
||||
image_data = b"fake_image_data"
|
||||
image = BytesIO(image_data)
|
||||
|
||||
image_edit_optional_params = {
|
||||
"n": 2,
|
||||
"response_format": "url",
|
||||
"strength": 0.5,
|
||||
"style": "photographic"
|
||||
}
|
||||
|
||||
litellm_params = GenericLiteLLMParams()
|
||||
headers = {}
|
||||
|
||||
data, files = self.config.transform_image_edit_request(
|
||||
model=self.model,
|
||||
prompt=self.prompt,
|
||||
image=image,
|
||||
image_edit_optional_request_params=image_edit_optional_params,
|
||||
litellm_params=litellm_params,
|
||||
headers=headers
|
||||
)
|
||||
|
||||
# Check that data contains the expected parameters
|
||||
assert data["model"] == self.model
|
||||
assert data["prompt"] == self.prompt
|
||||
assert data["strength"] == 0.5
|
||||
assert data["n"] == 2
|
||||
assert data["response_format"] == "url"
|
||||
assert data["style"] == "photographic"
|
||||
|
||||
# Check that image is not in data (should be in files)
|
||||
assert "image" not in data
|
||||
|
||||
# Check that files contains the image
|
||||
assert len(files) == 1
|
||||
assert files[0][0] == "image" # field name
|
||||
assert files[0][1][0] == "image.png" # filename (default for non-BufferedReader)
|
||||
assert files[0][1][1] == image # file object
|
||||
|
||||
def test_get_image_files_for_request_single_image(self):
|
||||
"""
|
||||
Test that _get_image_files_for_request correctly handles a single image.
|
||||
"""
|
||||
image_data = b"fake_image_data"
|
||||
image = BytesIO(image_data)
|
||||
|
||||
files = self.config._get_image_files_for_request(image=image)
|
||||
|
||||
assert len(files) == 1
|
||||
assert files[0][0] == "image"
|
||||
assert files[0][1][0] == "image.png" # Default filename for non-BufferedReader
|
||||
assert files[0][1][1] == image
|
||||
assert "image/png" in files[0][1][2]
|
||||
|
||||
def test_get_image_files_for_request_list_with_single_image(self):
|
||||
"""
|
||||
Test that _get_image_files_for_request correctly handles a list containing a single image
|
||||
(takes the first image for Recraft API).
|
||||
"""
|
||||
image_data = b"fake_image_data"
|
||||
image = BytesIO(image_data)
|
||||
|
||||
# Pass as list (OpenAI format)
|
||||
files = self.config._get_image_files_for_request(image=[image])
|
||||
|
||||
assert len(files) == 1
|
||||
assert files[0][0] == "image"
|
||||
assert files[0][1][0] == "image.png" # Default filename for non-BufferedReader
|
||||
assert files[0][1][1] == image
|
||||
|
||||
def test_get_image_files_for_request_buffered_reader(self):
|
||||
"""
|
||||
Test that _get_image_files_for_request correctly handles BufferedReader objects.
|
||||
"""
|
||||
# Create a mock BufferedReader
|
||||
mock_file = MagicMock(spec=BufferedReader)
|
||||
mock_file.name = "buffered_image.jpg"
|
||||
|
||||
files = self.config._get_image_files_for_request(image=mock_file)
|
||||
|
||||
assert len(files) == 1
|
||||
assert files[0][0] == "image"
|
||||
assert files[0][1][0] == "buffered_image.jpg"
|
||||
assert files[0][1][1] == mock_file
|
||||
|
||||
def test_get_image_files_for_request_no_image(self):
|
||||
"""
|
||||
Test that _get_image_files_for_request returns empty list when no image is provided.
|
||||
"""
|
||||
files = self.config._get_image_files_for_request(image=None)
|
||||
assert files == []
|
||||
|
||||
def test_transform_image_edit_response_success(self):
|
||||
"""
|
||||
Test that transform_image_edit_response correctly transforms a successful response.
|
||||
"""
|
||||
# Mock response data
|
||||
response_data = {
|
||||
"data": [
|
||||
{"url": "https://example.com/edited_image1.jpg", "b64_json": None},
|
||||
{"url": None, "b64_json": "base64encodeddata"}
|
||||
]
|
||||
}
|
||||
|
||||
# Create mock response
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = response_data
|
||||
|
||||
result = self.config.transform_image_edit_response(
|
||||
model=self.model,
|
||||
raw_response=mock_response,
|
||||
logging_obj=self.logging_obj
|
||||
)
|
||||
|
||||
assert isinstance(result, ImageResponse)
|
||||
assert len(result.data) == 2
|
||||
assert result.data[0].url == "https://example.com/edited_image1.jpg"
|
||||
assert result.data[0].b64_json is None
|
||||
assert result.data[1].url is None
|
||||
assert result.data[1].b64_json == "base64encodeddata"
|
||||
|
||||
def test_transform_image_edit_response_json_error(self):
|
||||
"""
|
||||
Test that transform_image_edit_response raises appropriate error when response JSON is invalid.
|
||||
"""
|
||||
# Create mock response that raises JSON decode error
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.side_effect = json.JSONDecodeError("Invalid JSON", "", 0)
|
||||
mock_response.status_code = 500
|
||||
mock_response.headers = {}
|
||||
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
self.config.transform_image_edit_response(
|
||||
model=self.model,
|
||||
raw_response=mock_response,
|
||||
logging_obj=self.logging_obj
|
||||
)
|
||||
|
||||
assert "Error transforming image edit response" in str(exc_info.value)
|
||||
@@ -0,0 +1,270 @@
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from typing import List, Optional
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../../../../..")
|
||||
) # Adds the parent directory to the system path
|
||||
|
||||
from litellm.llms.recraft.image_generation.transformation import (
|
||||
RecraftImageGenerationConfig,
|
||||
)
|
||||
from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams
|
||||
from litellm.types.utils import ImageObject, ImageResponse
|
||||
|
||||
|
||||
class TestRecraftImageGenerationTransformation:
|
||||
def setup_method(self):
|
||||
"""Set up test fixtures before each test method."""
|
||||
self.config = RecraftImageGenerationConfig()
|
||||
self.model = "recraft-v3"
|
||||
self.logging_obj = MagicMock()
|
||||
|
||||
|
||||
def test_map_openai_params_supported_params(self):
|
||||
"""Test that map_openai_params correctly maps supported parameters."""
|
||||
non_default_params = {
|
||||
"n": 2,
|
||||
"response_format": "url",
|
||||
"size": "1024x1024",
|
||||
"style": "photographic"
|
||||
}
|
||||
optional_params = {}
|
||||
|
||||
result = self.config.map_openai_params(
|
||||
non_default_params=non_default_params,
|
||||
optional_params=optional_params,
|
||||
model=self.model,
|
||||
drop_params=False
|
||||
)
|
||||
|
||||
assert result == non_default_params
|
||||
|
||||
def test_map_openai_params_unsupported_param_drop_true(self):
|
||||
"""Test that map_openai_params drops unsupported parameters when drop_params=True."""
|
||||
non_default_params = {
|
||||
"n": 2,
|
||||
"unsupported_param": "value"
|
||||
}
|
||||
optional_params = {}
|
||||
|
||||
result = self.config.map_openai_params(
|
||||
non_default_params=non_default_params,
|
||||
optional_params=optional_params,
|
||||
model=self.model,
|
||||
drop_params=True
|
||||
)
|
||||
|
||||
assert result == {"n": 2}
|
||||
assert "unsupported_param" not in result
|
||||
|
||||
def test_map_openai_params_unsupported_param_drop_false(self):
|
||||
"""Test that map_openai_params raises ValueError for unsupported parameters when drop_params=False."""
|
||||
non_default_params = {
|
||||
"n": 2,
|
||||
"unsupported_param": "value"
|
||||
}
|
||||
optional_params = {}
|
||||
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
self.config.map_openai_params(
|
||||
non_default_params=non_default_params,
|
||||
optional_params=optional_params,
|
||||
model=self.model,
|
||||
drop_params=False
|
||||
)
|
||||
|
||||
assert "unsupported_param" in str(exc_info.value)
|
||||
assert "is not supported for model" in str(exc_info.value)
|
||||
|
||||
@patch("litellm.llms.recraft.image_generation.transformation.get_secret_str")
|
||||
def test_get_complete_url_with_api_base(self, mock_get_secret):
|
||||
"""Test that get_complete_url returns correct URL when api_base is provided."""
|
||||
api_base = "https://custom.api.recraft.ai"
|
||||
|
||||
result = self.config.get_complete_url(
|
||||
api_base=api_base,
|
||||
api_key="test_key",
|
||||
model=self.model,
|
||||
optional_params={},
|
||||
litellm_params={}
|
||||
)
|
||||
|
||||
expected_url = f"{api_base}/{self.config.IMAGE_GENERATION_ENDPOINT}"
|
||||
assert result == expected_url
|
||||
mock_get_secret.assert_not_called()
|
||||
|
||||
@patch("litellm.llms.recraft.image_generation.transformation.get_secret_str")
|
||||
def test_get_complete_url_with_secret_base(self, mock_get_secret):
|
||||
"""Test that get_complete_url uses secret when api_base is None."""
|
||||
mock_get_secret.return_value = "https://secret.api.recraft.ai"
|
||||
|
||||
result = self.config.get_complete_url(
|
||||
api_base=None,
|
||||
api_key="test_key",
|
||||
model=self.model,
|
||||
optional_params={},
|
||||
litellm_params={}
|
||||
)
|
||||
|
||||
expected_url = f"https://secret.api.recraft.ai/{self.config.IMAGE_GENERATION_ENDPOINT}"
|
||||
assert result == expected_url
|
||||
mock_get_secret.assert_called_once_with("RECRAFT_API_BASE")
|
||||
|
||||
@patch("litellm.llms.recraft.image_generation.transformation.get_secret_str")
|
||||
def test_get_complete_url_with_default_base(self, mock_get_secret):
|
||||
"""Test that get_complete_url uses default base URL when no other options are available."""
|
||||
mock_get_secret.return_value = None
|
||||
|
||||
result = self.config.get_complete_url(
|
||||
api_base=None,
|
||||
api_key="test_key",
|
||||
model=self.model,
|
||||
optional_params={},
|
||||
litellm_params={}
|
||||
)
|
||||
|
||||
expected_url = f"{self.config.DEFAULT_BASE_URL}/{self.config.IMAGE_GENERATION_ENDPOINT}"
|
||||
assert result == expected_url
|
||||
|
||||
@patch("litellm.llms.recraft.image_generation.transformation.get_secret_str")
|
||||
def test_validate_environment_with_api_key(self, mock_get_secret):
|
||||
"""Test that validate_environment correctly sets authorization header when api_key is provided."""
|
||||
headers = {}
|
||||
api_key = "test_api_key"
|
||||
|
||||
result = self.config.validate_environment(
|
||||
headers=headers,
|
||||
model=self.model,
|
||||
messages=[],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
api_key=api_key
|
||||
)
|
||||
|
||||
assert result["Authorization"] == f"Bearer {api_key}"
|
||||
mock_get_secret.assert_not_called()
|
||||
|
||||
@patch("litellm.llms.recraft.image_generation.transformation.get_secret_str")
|
||||
def test_validate_environment_with_secret_key(self, mock_get_secret):
|
||||
"""Test that validate_environment uses secret API key when api_key is None."""
|
||||
mock_get_secret.return_value = "secret_api_key"
|
||||
headers = {}
|
||||
|
||||
result = self.config.validate_environment(
|
||||
headers=headers,
|
||||
model=self.model,
|
||||
messages=[],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
api_key=None
|
||||
)
|
||||
|
||||
assert result["Authorization"] == "Bearer secret_api_key"
|
||||
mock_get_secret.assert_called_once_with("RECRAFT_API_KEY")
|
||||
|
||||
@patch("litellm.llms.recraft.image_generation.transformation.get_secret_str")
|
||||
def test_validate_environment_no_api_key_raises_error(self, mock_get_secret):
|
||||
"""Test that validate_environment raises ValueError when no API key is available."""
|
||||
mock_get_secret.return_value = None
|
||||
headers = {}
|
||||
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
self.config.validate_environment(
|
||||
headers=headers,
|
||||
model=self.model,
|
||||
messages=[],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
api_key=None
|
||||
)
|
||||
|
||||
assert "RECRAFT_API_KEY is not set" in str(exc_info.value)
|
||||
|
||||
def test_transform_image_generation_request(self):
|
||||
"""Test that transform_image_generation_request correctly transforms request parameters."""
|
||||
prompt = "A beautiful sunset over mountains"
|
||||
optional_params = {
|
||||
"n": 2,
|
||||
"size": "1024x1024",
|
||||
"style": "photographic"
|
||||
}
|
||||
litellm_params = {}
|
||||
headers = {}
|
||||
|
||||
result = self.config.transform_image_generation_request(
|
||||
model=self.model,
|
||||
prompt=prompt,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
headers=headers
|
||||
)
|
||||
|
||||
assert result["prompt"] == prompt
|
||||
assert result["model"] == self.model
|
||||
assert result["n"] == 2
|
||||
assert result["size"] == "1024x1024"
|
||||
assert result["style"] == "photographic"
|
||||
|
||||
def test_transform_image_generation_response_success(self):
|
||||
"""Test that transform_image_generation_response correctly transforms successful response."""
|
||||
# Mock response data
|
||||
response_data = {
|
||||
"data": [
|
||||
{"url": "https://example.com/image1.jpg", "b64_json": None},
|
||||
{"url": None, "b64_json": "base64encodeddata"}
|
||||
]
|
||||
}
|
||||
|
||||
# Create mock response
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = response_data
|
||||
|
||||
# Create empty model response
|
||||
model_response = ImageResponse(data=[])
|
||||
|
||||
result = self.config.transform_image_generation_response(
|
||||
model=self.model,
|
||||
raw_response=mock_response,
|
||||
model_response=model_response,
|
||||
logging_obj=self.logging_obj,
|
||||
request_data={},
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
encoding=None
|
||||
)
|
||||
|
||||
assert len(result.data) == 2
|
||||
assert result.data[0].url == "https://example.com/image1.jpg"
|
||||
assert result.data[0].b64_json is None
|
||||
assert result.data[1].url is None
|
||||
assert result.data[1].b64_json == "base64encodeddata"
|
||||
|
||||
def test_transform_image_generation_response_json_error(self):
|
||||
"""Test that transform_image_generation_response raises error when response JSON is invalid."""
|
||||
# Create mock response that raises JSON decode error
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.side_effect = json.JSONDecodeError("Invalid JSON", "", 0)
|
||||
mock_response.status_code = 500
|
||||
mock_response.headers = {}
|
||||
|
||||
model_response = ImageResponse(data=[])
|
||||
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
self.config.transform_image_generation_response(
|
||||
model=self.model,
|
||||
raw_response=mock_response,
|
||||
model_response=model_response,
|
||||
logging_obj=self.logging_obj,
|
||||
request_data={},
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
encoding=None
|
||||
)
|
||||
|
||||
assert "Error transforming image generation response" in str(exc_info.value)
|
||||
Reference in New Issue
Block a user