113 lines
4.0 KiB
Python
113 lines
4.0 KiB
Python
import pytest
|
|
|
|
from app.integrations.openai_compatible import OpenAICompatibleGateway
|
|
|
|
|
|
def image_values(profile: str, model_id: str) -> dict:
|
|
return {
|
|
"model_pool": [
|
|
{
|
|
"id": "image",
|
|
"name": "测试生图",
|
|
"model_id": model_id,
|
|
"category": "multimodal",
|
|
"provider": "lingke",
|
|
"base_url": "https://api.lk888.ai",
|
|
"api_key": "secret",
|
|
"capabilities": ["image_generation"],
|
|
"models_path": "/v1/models",
|
|
"image_protocol": "aigc_media",
|
|
"image_parameter_profile": profile,
|
|
"image_generation_path": "/v1/media/generate",
|
|
"image_status_path": "/v1/media/status",
|
|
"enabled": True,
|
|
}
|
|
],
|
|
"image_routing_mode": "manual",
|
|
"image_model_id": "image",
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
("profile", "model_id", "options", "expected_params"),
|
|
[
|
|
("gpt_image_2", "gpt-image-2", {"size": "1024x1024", "quality": "low"}, {"size": "1024x1024", "quality": "low"}),
|
|
("nano_banana_pro", "gemini-3-pro-image-preview", {"aspect_ratio": "1:1", "size": "1K"}, {"aspectRatio": "1:1", "imageSize": "1K"}),
|
|
("seedream_5_pro", "doubao-seedream-5-0-pro-260628", {"aspectRatio": "1:1", "imageSize": "1K"}, {"aspect_ratio": "1:1", "size": "1K"}),
|
|
],
|
|
)
|
|
async def test_aigc_media_profiles_send_documented_parameter_names(
|
|
monkeypatch,
|
|
profile: str,
|
|
model_id: str,
|
|
options: dict,
|
|
expected_params: dict,
|
|
) -> None:
|
|
requests: list[dict] = []
|
|
|
|
class FakeResponse:
|
|
def raise_for_status(self) -> None:
|
|
return None
|
|
|
|
def json(self) -> dict:
|
|
return {"data": [{"url": "https://cdn.example.com/test.png"}]}
|
|
|
|
class FakeClient:
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, *args) -> None:
|
|
return None
|
|
|
|
async def post(self, url, **kwargs) -> FakeResponse:
|
|
requests.append({"url": url, **kwargs})
|
|
return FakeResponse()
|
|
|
|
monkeypatch.setattr("app.integrations.openai_compatible.httpx.AsyncClient", lambda **kwargs: FakeClient())
|
|
|
|
result = await OpenAICompatibleGateway(image_values(profile, model_id)).generate_image("测试", **options)
|
|
|
|
assert result["data"][0]["url"].endswith("test.png")
|
|
assert requests[0]["url"] == "https://api.lk888.ai/v1/media/generate"
|
|
assert requests[0]["json"] == {"model": model_id, "prompt": "测试", "params": expected_params}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aigc_media_polls_task_until_result_url(monkeypatch) -> None:
|
|
class FakeResponse:
|
|
def __init__(self, payload: dict) -> None:
|
|
self.payload = payload
|
|
|
|
def raise_for_status(self) -> None:
|
|
return None
|
|
|
|
def json(self) -> dict:
|
|
return self.payload
|
|
|
|
class FakeClient:
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, *args) -> None:
|
|
return None
|
|
|
|
async def post(self, *args, **kwargs) -> FakeResponse:
|
|
return FakeResponse({"task_id": 123})
|
|
|
|
async def get(self, url, **kwargs) -> FakeResponse:
|
|
assert url == "https://api.lk888.ai/v1/media/status"
|
|
assert kwargs["params"] == {"task_id": 123}
|
|
return FakeResponse({"task_id": 123, "state": "success", "is_final": True, "result_url": "https://cdn.example.com/final.png"})
|
|
|
|
async def no_sleep(*args) -> None:
|
|
return None
|
|
|
|
monkeypatch.setattr("app.integrations.openai_compatible.httpx.AsyncClient", lambda **kwargs: FakeClient())
|
|
monkeypatch.setattr("app.integrations.openai_compatible.asyncio.sleep", no_sleep)
|
|
|
|
result = await OpenAICompatibleGateway(image_values("seedream_5_pro", "doubao-seedream-5-0-pro-260628")).generate_image("测试")
|
|
|
|
assert result["state"] == "success"
|
|
assert result["data"] == [{"url": "https://cdn.example.com/final.png"}]
|