Grayscale/nanobot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'origin/main' into codex/review-pr-3929

Browse Source
This commit is contained in:
Xubin Ren
2026-05-22 22:15:46 +08:00
parent 3f789bd9f9 055c9be359
commit 143224e25a
72 changed files with 9551 additions and 829 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/providers/test_image_generation.py
+496
View File
@@ -9,10 +9,13 @@ import pytest
from nanobot.providers.image_generation import (
AIHubMixImageGenerationClient,
CodexImageGenerationClient,
GeminiImageGenerationClient,
GeneratedImageResponse,
ImageGenerationError,
MiniMaxImageGenerationClient,
OllamaImageGenerationClient,
OpenAIImageGenerationClient,
OpenRouterImageGenerationClient,
StepFunImageGenerationClient,
)
@@ -36,12 +39,14 @@ class FakeResponse:
payload: dict[str, Any],
status_code: int = 200,
content: bytes = b"",
sse_lines: list[str] | None = None,
) -> None:
self._payload = payload
self.status_code = status_code
self.text = str(payload)
self.content = content
self.request = httpx.Request("POST", "https://openrouter.ai/api/v1/chat/completions")
self._sse_lines = sse_lines
def json(self) -> dict[str, Any]:
return self._payload
@@ -51,6 +56,15 @@ class FakeResponse:
response = httpx.Response(self.status_code, request=self.request, text=self.text)
raise httpx.HTTPStatusError("failed", request=self.request, response=response)
async def aiter_lines(self):
if self._sse_lines is not None:
for line in self._sse_lines:
yield line
return
# Fallback: treat response text as SSE lines
for line in self.text.split("\n"):
yield line
class FakeClient:
def __init__(self, response: FakeResponse) -> None:
@@ -133,6 +147,54 @@ async def test_openrouter_image_generation_requires_api_key() -> None:
await client.generate(prompt="draw", model="model")
@pytest.mark.asyncio
async def test_ollama_image_generation_payload_and_response() -> None:
raw_b64 = PNG_DATA_URL.removeprefix("data:image/png;base64,")
fake = FakeClient(FakeResponse({"image": raw_b64}))
client = OllamaImageGenerationClient(
api_key="ollama-test",
api_base="http://localhost:11434/v1/",
extra_headers={"X-Test": "1"},
extra_body={"seed": 123},
client=fake, # type: ignore[arg-type]
)
response = await client.generate(
prompt="a sunset",
model="x/z-image-turbo",
aspect_ratio="16:9",
image_size="1K",
)
assert response.images == [PNG_DATA_URL]
assert response.content == ""
call = fake.calls[0]
assert call["url"] == "http://localhost:11434/api/generate"
assert call["headers"]["Authorization"] == "Bearer ollama-test"
assert call["headers"]["X-Test"] == "1"
body = call["json"]
assert body["model"] == "x/z-image-turbo"
assert body["prompt"] == "a sunset"
assert body["width"] == 1024
assert body["height"] == 576
assert body["steps"] == 0
assert body["stream"] is False
assert body["seed"] == 123
@pytest.mark.asyncio
async def test_ollama_image_generation_rejects_reference_images() -> None:
client = OllamaImageGenerationClient(api_key=None)
with pytest.raises(ImageGenerationError, match="reference images"):
await client.generate(
prompt="edit this",
model="x/z-image-turbo",
reference_images=["ref.png"],
)
@pytest.mark.asyncio
async def test_aihubmix_image_generation_payload_and_response() -> None:
raw_b64 = PNG_DATA_URL.removeprefix("data:image/png;base64,")
@@ -531,3 +593,437 @@ async def test_stepfun_no_images_raises() -> None:
with pytest.raises(ImageGenerationError, match="returned no images"):
await client.generate(prompt="draw", model="step-image-edit-2")
# ---------------------------------------------------------------------------
# OpenAI
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_openai_payload_and_response() -> None:
fake = FakeClient(FakeResponse({"data": [{"b64_json": RAW_B64}]}))
client = OpenAIImageGenerationClient(
api_key="sk-openai-test",
api_base="https://api.openai.com/v1",
extra_headers={"X-Test": "1"},
client=fake, # type: ignore[arg-type]
)
response = await client.generate(
prompt="a cat on the moon",
model="dall-e-3",
aspect_ratio="16:9",
)
assert response.images == [PNG_DATA_URL]
call = fake.calls[0]
assert call["url"] == "https://api.openai.com/v1/images/generations"
assert call["headers"]["Authorization"] == "Bearer sk-openai-test"
assert call["headers"]["X-Test"] == "1"
body = call["json"]
assert body["model"] == "dall-e-3"
assert body["prompt"] == "a cat on the moon"
assert body["response_format"] == "b64_json"
assert body["n"] == 1
assert body["size"] == "1792x1024"
@pytest.mark.asyncio
async def test_openai_b64_json_response_uses_detected_mime() -> None:
raw_b64 = base64.b64encode(JPEG_BYTES).decode("ascii")
fake = FakeClient(FakeResponse({"data": [{"b64_json": raw_b64}]}))
client = OpenAIImageGenerationClient(
api_key="sk-openai-test",
client=fake, # type: ignore[arg-type]
)
response = await client.generate(prompt="draw", model="dall-e-3")
assert response.images == [f"data:image/jpeg;base64,{raw_b64}"]
@pytest.mark.asyncio
async def test_openai_url_download_fallback() -> None:
fake = FakeClient(FakeResponse({"data": [{"url": "https://cdn.example/image.png"}]}))
fake.get_response = FakeResponse({}, content=PNG_BYTES)
client = OpenAIImageGenerationClient(
api_key="sk-openai-test",
client=fake, # type: ignore[arg-type]
)
response = await client.generate(prompt="draw", model="dall-e-3")
assert response.images[0].startswith("data:image/png;base64,")
assert fake.get_calls[0]["url"] == "https://cdn.example/image.png"
@pytest.mark.asyncio
async def test_openai_multiple_images() -> None:
fake = FakeClient(FakeResponse({
"data": [
{"b64_json": RAW_B64},
{"b64_json": RAW_B64},
]
}))
client = OpenAIImageGenerationClient(
api_key="sk-openai-test",
client=fake, # type: ignore[arg-type]
)
response = await client.generate(prompt="draw", model="dall-e-3")
assert len(response.images) == 2
assert response.images == [PNG_DATA_URL, PNG_DATA_URL]
@pytest.mark.asyncio
async def test_openai_aspect_ratio_to_size() -> None:
fake = FakeClient(FakeResponse({"data": [{"b64_json": RAW_B64}]}))
client = OpenAIImageGenerationClient(
api_key="sk-openai-test",
client=fake, # type: ignore[arg-type]
)
await client.generate(prompt="draw", model="dall-e-3", aspect_ratio="1:1")
assert fake.calls[0]["json"]["size"] == "1024x1024"
@pytest.mark.asyncio
async def test_openai_dalle3_uses_supported_orientation_sizes() -> None:
fake = FakeClient(FakeResponse({"data": [{"b64_json": RAW_B64}]}))
client = OpenAIImageGenerationClient(
api_key="sk-openai-test",
client=fake, # type: ignore[arg-type]
)
await client.generate(prompt="draw", model="dall-e-3", aspect_ratio="3:4")
await client.generate(prompt="draw", model="dall-e-3", aspect_ratio="4:3")
assert fake.calls[0]["json"]["size"] == "1024x1792"
assert fake.calls[1]["json"]["size"] == "1792x1024"
@pytest.mark.asyncio
async def test_openai_dalle2_uses_square_size_for_non_square_ratios() -> None:
fake = FakeClient(FakeResponse({"data": [{"b64_json": RAW_B64}]}))
client = OpenAIImageGenerationClient(
api_key="sk-openai-test",
client=fake, # type: ignore[arg-type]
)
await client.generate(prompt="draw", model="dall-e-2", aspect_ratio="16:9")
assert fake.calls[0]["json"]["size"] == "1024x1024"
@pytest.mark.asyncio
async def test_openai_gpt_image_uses_supported_landscape_size() -> None:
fake = FakeClient(FakeResponse({"data": [{"b64_json": RAW_B64}]}))
client = OpenAIImageGenerationClient(
api_key="sk-openai-test",
client=fake, # type: ignore[arg-type]
)
await client.generate(prompt="draw", model="gpt-image-1", aspect_ratio="16:9")
assert fake.calls[0]["json"]["size"] == "1536x1024"
@pytest.mark.asyncio
async def test_openai_gpt_image_uses_supported_orientation_sizes() -> None:
fake = FakeClient(FakeResponse({"data": [{"b64_json": RAW_B64}]}))
client = OpenAIImageGenerationClient(
api_key="sk-openai-test",
client=fake, # type: ignore[arg-type]
)
await client.generate(prompt="draw", model="gpt-image-1", aspect_ratio="3:4")
await client.generate(prompt="draw", model="gpt-image-1", aspect_ratio="4:3")
assert fake.calls[0]["json"]["size"] == "1024x1536"
assert fake.calls[1]["json"]["size"] == "1536x1024"
@pytest.mark.asyncio
async def test_openai_default_size_when_no_aspect_ratio() -> None:
fake = FakeClient(FakeResponse({"data": [{"b64_json": RAW_B64}]}))
client = OpenAIImageGenerationClient(
api_key="sk-openai-test",
client=fake, # type: ignore[arg-type]
)
await client.generate(prompt="draw", model="dall-e-3")
body = fake.calls[0]["json"]
assert body["size"] == "1024x1024"
@pytest.mark.asyncio
async def test_openai_ignores_explicit_size_unsupported_by_model_family() -> None:
fake = FakeClient(FakeResponse({"data": [{"b64_json": RAW_B64}]}))
client = OpenAIImageGenerationClient(
api_key="sk-openai-test",
client=fake, # type: ignore[arg-type]
)
await client.generate(
prompt="draw",
model="dall-e-3",
aspect_ratio="16:9",
image_size="1536x1024",
)
body = fake.calls[0]["json"]
assert body["size"] == "1792x1024"
@pytest.mark.asyncio
async def test_openai_uses_explicit_image_size() -> None:
fake = FakeClient(FakeResponse({"data": [{"b64_json": RAW_B64}]}))
client = OpenAIImageGenerationClient(
api_key="sk-openai-test",
client=fake, # type: ignore[arg-type]
)
await client.generate(
prompt="draw",
model="dall-e-3",
aspect_ratio="16:9",
image_size="1024x1024",
)
body = fake.calls[0]["json"]
assert body["size"] == "1024x1024"
@pytest.mark.asyncio
async def test_openai_requires_api_key() -> None:
client = OpenAIImageGenerationClient(api_key=None)
with pytest.raises(ImageGenerationError, match="API key"):
await client.generate(prompt="draw", model="dall-e-3")
# ---------------------------------------------------------------------------
# OpenAI Codex (Responses API)
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_codex_payload_and_response(monkeypatch) -> None:
import sys
from dataclasses import dataclass
from types import SimpleNamespace
@dataclass
class FakeToken:
account_id: str = "acct-123"
access: str = "oauth-token"
async def fake_to_thread(fn, *args, **kwargs):
return fn(*args, **kwargs)
monkeypatch.setattr("asyncio.to_thread", fake_to_thread)
fake_oauth = SimpleNamespace(get_token=lambda: FakeToken())
monkeypatch.setitem(sys.modules, "oauth_cli_kit", fake_oauth)
sse_lines = [
'data: {"type":"response.output_item.added","item":{"id":"ig_1","type":"image_generation_call","status":"in_progress"}}',
"",
f'data: {{"type":"response.output_item.done","item":{{"id":"ig_1","type":"image_generation_call","result":"{PNG_DATA_URL}","status":"completed"}}}}',
"",
'data: [DONE]',
"",
]
fake = FakeClient(FakeResponse({}, sse_lines=sse_lines))
client = CodexImageGenerationClient(
api_key=None,
api_base="https://chatgpt.com/backend-api",
extra_headers={"X-Test": "1"},
client=fake, # type: ignore[arg-type]
)
response = await client.generate(
prompt="draw a cat",
model="gpt-5.4",
)
assert response.images == [PNG_DATA_URL]
assert response.content == ""
call = fake.calls[0]
assert call["url"] == "https://chatgpt.com/backend-api/codex/responses"
assert call["headers"]["Authorization"] == "Bearer oauth-token"
assert call["headers"]["chatgpt-account-id"] == "acct-123"
assert call["headers"]["OpenAI-Beta"] == "responses=experimental"
assert call["headers"]["X-Test"] == "1"
body = call["json"]
assert body["model"] == "gpt-5.4"
assert body["instructions"] == "Generate an image based on the user's request."
assert body["input"] == [{"role": "user", "content": "draw a cat"}]
assert body["tools"] == [{"type": "image_generation"}]
assert body["tool_choice"] == "auto"
assert body["store"] is False
assert body["stream"] is True
@pytest.mark.asyncio
async def test_codex_strips_model_prefix(monkeypatch) -> None:
import sys
from dataclasses import dataclass
from types import SimpleNamespace
@dataclass
class FakeToken:
account_id: str = "acct-123"
access: str = "oauth-token"
async def fake_to_thread(fn, *args, **kwargs):
return fn(*args, **kwargs)
monkeypatch.setattr("asyncio.to_thread", fake_to_thread)
fake_oauth = SimpleNamespace(get_token=lambda: FakeToken())
monkeypatch.setitem(sys.modules, "oauth_cli_kit", fake_oauth)
fake = FakeClient(FakeResponse({}, sse_lines=[
f'data: {{"type":"response.output_item.done","item":{{"type":"image_generation_call","result":"{PNG_DATA_URL}"}}}}',
"",
'data: [DONE]',
"",
]))
client = CodexImageGenerationClient(
api_key=None, client=fake # type: ignore[arg-type]
)
await client.generate(prompt="draw", model="openai-codex/gpt-5.4")
assert fake.calls[0]["json"]["model"] == "gpt-5.4"
@pytest.mark.asyncio
async def test_codex_requires_oauth(monkeypatch) -> None:
async def fake_to_thread(fn, *args, **kwargs):
raise RuntimeError("no token")
monkeypatch.setattr("asyncio.to_thread", fake_to_thread)
client = CodexImageGenerationClient(api_key=None)
with pytest.raises(ImageGenerationError, match="OAuth token"):
await client.generate(prompt="draw", model="gpt-5.4")
@pytest.mark.asyncio
async def test_codex_no_images_raises(monkeypatch) -> None:
import sys
from dataclasses import dataclass
from types import SimpleNamespace
@dataclass
class FakeToken:
account_id: str = "acct-123"
access: str = "oauth-token"
async def fake_to_thread(fn, *args, **kwargs):
return fn(*args, **kwargs)
monkeypatch.setattr("asyncio.to_thread", fake_to_thread)
fake_oauth = SimpleNamespace(get_token=lambda: FakeToken())
monkeypatch.setitem(sys.modules, "oauth_cli_kit", fake_oauth)
fake = FakeClient(FakeResponse({}, sse_lines=[
'data: {"type":"response.completed","response":{"status":"completed"}}',
"",
'data: [DONE]',
"",
]))
client = CodexImageGenerationClient(
api_key=None, client=fake # type: ignore[arg-type]
)
with pytest.raises(ImageGenerationError, match="returned no images"):
await client.generate(prompt="draw", model="gpt-5.4")
@pytest.mark.asyncio
async def test_codex_extracts_text_content(monkeypatch) -> None:
import sys
from dataclasses import dataclass
from types import SimpleNamespace
@dataclass
class FakeToken:
account_id: str = "acct-123"
access: str = "oauth-token"
async def fake_to_thread(fn, *args, **kwargs):
return fn(*args, **kwargs)
monkeypatch.setattr("asyncio.to_thread", fake_to_thread)
fake_oauth = SimpleNamespace(get_token=lambda: FakeToken())
monkeypatch.setitem(sys.modules, "oauth_cli_kit", fake_oauth)
fake = FakeClient(FakeResponse({}, sse_lines=[
'data: {"type":"response.output_text.delta","delta":"Here "}',
"",
'data: {"type":"response.output_text.delta","delta":"is your cat image."}',
"",
f'data: {{"type":"response.output_item.done","item":{{"type":"image_generation_call","result":"{PNG_DATA_URL}"}}}}',
"",
'data: [DONE]',
"",
]))
client = CodexImageGenerationClient(
api_key=None, client=fake # type: ignore[arg-type]
)
response = await client.generate(prompt="draw a cat", model="gpt-5.4")
assert response.images == [PNG_DATA_URL]
assert response.content == "Here is your cat image."
@pytest.mark.asyncio
async def test_codex_json_result_format(monkeypatch) -> None:
"""image_generation_call result can be a dict with image_url key."""
import sys
from dataclasses import dataclass
from types import SimpleNamespace
@dataclass
class FakeToken:
account_id: str = "acct-123"
access: str = "oauth-token"
async def fake_to_thread(fn, *args, **kwargs):
return fn(*args, **kwargs)
monkeypatch.setattr("asyncio.to_thread", fake_to_thread)
fake_oauth = SimpleNamespace(get_token=lambda: FakeToken())
monkeypatch.setitem(sys.modules, "oauth_cli_kit", fake_oauth)
fake = FakeClient(FakeResponse({}, sse_lines=[
f'data: {{"type":"response.output_item.done","item":{{"type":"image_generation_call","result":{{"image_url":"{PNG_DATA_URL}"}}}}}}',
"",
'data: [DONE]',
"",
]))
client = CodexImageGenerationClient(
api_key=None, client=fake # type: ignore[arg-type]
)
response = await client.generate(prompt="draw", model="gpt-5.4")
assert response.images == [PNG_DATA_URL]
@pytest.mark.asyncio
async def test_openai_no_images_raises() -> None:
fake = FakeClient(FakeResponse({"data": []}))
client = OpenAIImageGenerationClient(
api_key="sk-openai-test",
client=fake, # type: ignore[arg-type]
)
with pytest.raises(ImageGenerationError, match="returned no images"):
await client.generate(prompt="draw", model="dall-e-3")