wehub-resource-sync
216 lines
7.6 KiB
Python
216 lines
7.6 KiB
Python
"""Tests for the image_generation tool."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
from openharness.config.settings import ImageGenerationConfig
|
|
from openharness.tools.base import ToolExecutionContext
|
|
from openharness.tools.image_generation_tool import ImageGenerationTool, ImageGenerationToolInput
|
|
|
|
|
|
class _FakeStreamResponse:
|
|
def __init__(self, *, status_code: int = 200, lines: list[str] | None = None, body: str = "") -> None:
|
|
self.status_code = status_code
|
|
self._lines = lines or []
|
|
self._body = body.encode("utf-8")
|
|
|
|
async def __aenter__(self) -> "_FakeStreamResponse":
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc, tb) -> None:
|
|
return None
|
|
|
|
async def aread(self) -> bytes:
|
|
return self._body
|
|
|
|
async def aiter_lines(self):
|
|
for line in self._lines:
|
|
yield line
|
|
|
|
|
|
class _FakeAsyncClient:
|
|
def __init__(self, response: _FakeStreamResponse, sink: dict[str, Any]) -> None:
|
|
self._response = response
|
|
self._sink = sink
|
|
|
|
async def __aenter__(self) -> "_FakeAsyncClient":
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc, tb) -> None:
|
|
return None
|
|
|
|
def stream(self, method: str, url: str, *, headers: dict[str, str], json: dict[str, Any]):
|
|
self._sink["method"] = method
|
|
self._sink["url"] = url
|
|
self._sink["headers"] = headers
|
|
self._sink["json"] = json
|
|
return self._response
|
|
|
|
|
|
def _b64url(data: dict[str, object]) -> str:
|
|
raw = json.dumps(data, separators=(",", ":")).encode("utf-8")
|
|
return base64.urlsafe_b64encode(raw).decode("ascii").rstrip("=")
|
|
|
|
|
|
def _fake_codex_token() -> str:
|
|
payload = {"https://api.openai.com/auth": {"chatgpt_account_id": "acct_test"}}
|
|
return f"{_b64url({'alg': 'none', 'typ': 'JWT'})}.{_b64url(payload)}.sig"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_requires_api_key_for_openai(tmp_path: Path) -> None:
|
|
tool = ImageGenerationTool()
|
|
result = await tool.execute(
|
|
ImageGenerationToolInput(prompt="a cat", provider="openai"),
|
|
ToolExecutionContext(cwd=tmp_path, metadata={"image_generation_config": {}}),
|
|
)
|
|
assert result.is_error
|
|
assert "API key is not configured" in result.output
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_generate_writes_file(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
|
|
image_bytes = b"fake-png"
|
|
image_b64 = base64.b64encode(image_bytes).decode("ascii")
|
|
|
|
async def fake_generate_images(arguments, model, api_key, base_url):
|
|
assert arguments.prompt == "a cat"
|
|
assert model == "gpt-image-2"
|
|
assert api_key == "test-key"
|
|
assert base_url == ""
|
|
return [image_b64]
|
|
|
|
monkeypatch.setattr(ImageGenerationTool, "_generate_images", staticmethod(fake_generate_images))
|
|
|
|
tool = ImageGenerationTool()
|
|
result = await tool.execute(
|
|
ImageGenerationToolInput(prompt="a cat", output_path="assets/cat.png", provider="openai"),
|
|
ToolExecutionContext(
|
|
cwd=tmp_path,
|
|
metadata={"image_generation_config": {"api_key": "test-key", "model": "gpt-image-2"}},
|
|
),
|
|
)
|
|
|
|
out = tmp_path / "assets" / "cat.png"
|
|
assert not result.is_error
|
|
assert out.read_bytes() == image_bytes
|
|
assert result.metadata["paths"] == [str(out)]
|
|
assert result.metadata["provider"] == "openai"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_codex_hosted_generation(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
|
|
image_bytes = b"codex-png"
|
|
image_b64 = base64.b64encode(image_bytes).decode("ascii")
|
|
sink: dict[str, Any] = {}
|
|
response = _FakeStreamResponse(
|
|
lines=[
|
|
'data: {"type":"response.output_item.done","item":{"id":"ig_1","type":"image_generation_call","status":"completed","revised_prompt":"blue icon","result":"%s"}}' % image_b64,
|
|
"",
|
|
'data: {"type":"response.completed","response":{"status":"completed"}}',
|
|
"",
|
|
]
|
|
)
|
|
monkeypatch.setattr(
|
|
"openharness.tools.image_generation_tool.httpx.AsyncClient",
|
|
lambda *args, **kwargs: _FakeAsyncClient(response, sink),
|
|
)
|
|
|
|
tool = ImageGenerationTool()
|
|
result = await tool.execute(
|
|
ImageGenerationToolInput(prompt="a blue icon", output_path="icon.png", provider="codex"),
|
|
ToolExecutionContext(
|
|
cwd=tmp_path,
|
|
metadata={
|
|
"image_generation_config": {
|
|
"codex_auth_token": _fake_codex_token(),
|
|
"codex_model": "gpt-5.4",
|
|
}
|
|
},
|
|
),
|
|
)
|
|
|
|
out = tmp_path / "icon.png"
|
|
assert not result.is_error
|
|
assert out.read_bytes() == image_bytes
|
|
assert result.metadata["provider"] == "codex"
|
|
assert result.metadata["revised_prompt"] == "blue icon"
|
|
assert sink["url"].endswith("/codex/responses")
|
|
assert sink["json"]["tools"] == [{"type": "image_generation", "output_format": "png"}]
|
|
assert sink["json"]["tool_choice"] == "auto"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_auto_provider_prefers_codex_when_token_available(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
|
|
image_b64 = base64.b64encode(b"codex").decode("ascii")
|
|
|
|
async def fake_codex(self, arguments, config):
|
|
return [image_b64], None
|
|
|
|
monkeypatch.setattr(ImageGenerationTool, "_generate_with_codex", fake_codex)
|
|
|
|
tool = ImageGenerationTool()
|
|
result = await tool.execute(
|
|
ImageGenerationToolInput(prompt="a cat", output_path="cat.png"),
|
|
ToolExecutionContext(
|
|
cwd=tmp_path,
|
|
metadata={"image_generation_config": {"codex_auth_token": _fake_codex_token()}},
|
|
),
|
|
)
|
|
|
|
assert not result.is_error
|
|
assert result.metadata["provider"] == "codex"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_refuses_overwrite_by_default(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
|
|
out = tmp_path / "image.png"
|
|
out.write_bytes(b"existing")
|
|
image_b64 = base64.b64encode(b"new").decode("ascii")
|
|
|
|
async def fake_generate_images(arguments, model, api_key, base_url):
|
|
return [image_b64]
|
|
|
|
monkeypatch.setattr(ImageGenerationTool, "_generate_images", staticmethod(fake_generate_images))
|
|
|
|
tool = ImageGenerationTool()
|
|
result = await tool.execute(
|
|
ImageGenerationToolInput(prompt="a cat", output_path=str(out), provider="openai"),
|
|
ToolExecutionContext(cwd=tmp_path, metadata={"image_generation_config": {"api_key": "test"}}),
|
|
)
|
|
|
|
assert result.is_error
|
|
assert "output already exists" in result.output
|
|
assert out.read_bytes() == b"existing"
|
|
|
|
|
|
def test_resolve_output_paths_multiple(tmp_path: Path) -> None:
|
|
paths = ImageGenerationTool._resolve_output_paths(
|
|
ImageGenerationToolInput(output_path="hero.png", n=2),
|
|
tmp_path,
|
|
)
|
|
assert paths == [tmp_path / "hero-1.png", tmp_path / "hero-2.png"]
|
|
|
|
|
|
def test_image_generation_config_from_env(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("OPENHARNESS_IMAGE_GENERATION_PROVIDER", "openai")
|
|
monkeypatch.setenv("OPENHARNESS_IMAGE_GENERATION_MODEL", "gpt-image-1")
|
|
monkeypatch.setenv("OPENHARNESS_IMAGE_GENERATION_API_KEY", "sk-test")
|
|
monkeypatch.setenv("OPENHARNESS_IMAGE_GENERATION_BASE_URL", "https://example.test/v1")
|
|
monkeypatch.setenv("OPENHARNESS_IMAGE_GENERATION_CODEX_MODEL", "gpt-5.4")
|
|
|
|
cfg = ImageGenerationConfig.from_env()
|
|
|
|
assert cfg.provider == "openai"
|
|
assert cfg.model == "gpt-image-1"
|
|
assert cfg.api_key == "sk-test"
|
|
assert cfg.base_url == "https://example.test/v1"
|
|
assert cfg.codex_model == "gpt-5.4"
|
|
assert cfg.is_configured
|