Files
zhuangxiu/services/api/tests/test_image_gateway.py
T

283 lines
9.2 KiB
Python

import pytest
from app.integrations.openai_compatible import OpenAICompatibleGateway
from app.integrations.openai_compatible import (
extract_task_id,
safe_response_diagnostic,
unwrap_media_status,
)
from app.runtime_settings import (
RuntimeSettingsTestResult,
_verify_generated_image,
test_model_pool_item as run_model_pool_test,
)
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"}]
@pytest.mark.asyncio
async def test_aigc_media_connection_accepts_model_list_omission(monkeypatch) -> None:
class FakeResponse:
def raise_for_status(self) -> None:
return None
def json(self) -> dict:
return {"data": [{"id": "language-model-only"}]}
class FakeClient:
async def __aenter__(self):
return self
async def __aexit__(self, *args) -> None:
return None
async def get(self, *args, **kwargs) -> FakeResponse:
return FakeResponse()
monkeypatch.setattr("app.runtime_settings.httpx.AsyncClient", lambda **kwargs: FakeClient())
result = await run_model_pool_test(
image_values("gpt_image_2", "gpt-image-2"),
"image",
)
assert result.ok is True
assert result.details["id_advertised"] is False
assert result.details["verification"] == "connectivity_only"
assert "最终可用性以“试生成”为准" in result.message
@pytest.mark.asyncio
async def test_generation_endpoint_records_definitive_model_result(monkeypatch) -> None:
from app.api.settings import test_model_generation
class FakeStore:
def __init__(self) -> None:
self.recorded: list[RuntimeSettingsTestResult] = []
def merged_values(self) -> dict:
return image_values("gpt_image_2", "gpt-image-2")
def record_test(self, result: RuntimeSettingsTestResult, values: dict) -> None:
self.recorded.append(result)
async def fake_generation(values: dict, model_id: str) -> RuntimeSettingsTestResult:
return RuntimeSettingsTestResult(
target=f"generation:{model_id}",
ok=True,
message="端到端可用",
)
monkeypatch.setattr("app.api.settings.test_image_generation", fake_generation)
store = FakeStore()
result = await test_model_generation("image", store=store)
assert result.ok is True
assert [item.target for item in store.recorded] == ["generation:image", "model:image"]
assert store.recorded[1].ok is True
assert "真实生图端到端验证" in store.recorded[1].message
@pytest.mark.asyncio
async def test_generated_image_url_must_download_as_real_image(monkeypatch) -> None:
png_bytes = b"\x89PNG\r\n\x1a\n" + b"x" * 2048
class FakeResponse:
content = png_bytes
headers = {"content-type": "image/png"}
def raise_for_status(self) -> None:
return None
class FakeClient:
async def __aenter__(self):
return self
async def __aexit__(self, *args) -> None:
return None
async def get(self, url: str) -> FakeResponse:
assert url == "https://cdn.example.com/verified.png"
return FakeResponse()
monkeypatch.setattr("app.runtime_settings.httpx.AsyncClient", lambda **kwargs: FakeClient())
details = await _verify_generated_image(("url", "https://cdn.example.com/verified.png"))
assert details["image_format"] == "png"
assert details["byte_size"] == len(png_bytes)
@pytest.mark.asyncio
async def test_generated_image_url_rejects_html_placeholder(monkeypatch) -> None:
class FakeResponse:
content = b"<html><body>not an image</body></html>" + b"x" * 2048
headers = {"content-type": "text/html"}
def raise_for_status(self) -> None:
return None
class FakeClient:
async def __aenter__(self):
return self
async def __aexit__(self, *args) -> None:
return None
async def get(self, url: str) -> FakeResponse:
return FakeResponse()
monkeypatch.setattr("app.runtime_settings.httpx.AsyncClient", lambda **kwargs: FakeClient())
with pytest.raises(ValueError, match="不是可识别的图片"):
await _verify_generated_image(("url", "https://cdn.example.com/not-image"))
def test_response_diagnostic_keeps_errors_but_redacts_payloads() -> None:
diagnostic = safe_response_diagnostic(
{
"code": 402,
"error": {"message": "insufficient balance"},
"data": [{"b64_json": "secret-image-bytes", "url": "https://private.example"}],
}
)
rendered = str(diagnostic)
assert "insufficient balance" in rendered
assert "secret-image-bytes" not in rendered
assert "https://private.example" not in rendered
def test_extract_task_id_accepts_aggregator_nested_response() -> None:
payload = {
"code": 200,
"data": {"task_id": 91584074, "task_ids": [91584074]},
"msg": "Task created successfully",
}
assert extract_task_id(payload) == 91584074
def test_unwrap_media_status_accepts_nested_data() -> None:
payload = {
"code": 200,
"data": {
"state": "success",
"is_final": True,
"result_url": "https://cdn.example.com/result.png",
},
}
status = unwrap_media_status(payload)
assert status["is_final"] is True
assert status["result_url"].endswith("result.png")