mirror of
https://github.com/ggozad/oterm.git
synced 2026-10-11 09:43:21 +02:00
92 lines
3.2 KiB
Python
92 lines
3.2 KiB
Python
import base64
|
|
|
|
import pytest
|
|
from ollama import AsyncClient, RequestError, ResponseError
|
|
from ollama._types import GenerateResponse
|
|
from pydantic_ai.exceptions import ModelRetry
|
|
from pydantic_ai.messages import BinaryImage
|
|
|
|
from oterm.config import envConfig
|
|
from oterm.log import log_lines
|
|
from oterm.tools.ollama_image import generate_image
|
|
|
|
|
|
def _patch_generate(monkeypatch, response: GenerateResponse, capture: dict):
|
|
async def fake_generate(self, **kwargs):
|
|
capture.update(kwargs)
|
|
return response
|
|
|
|
monkeypatch.setattr(AsyncClient, "generate", fake_generate)
|
|
|
|
|
|
def _patch_generate_raises(monkeypatch, exc: Exception):
|
|
async def fake_generate(self, **kwargs):
|
|
raise exc
|
|
|
|
monkeypatch.setattr(AsyncClient, "generate", fake_generate)
|
|
|
|
|
|
class TestGenerateImage:
|
|
async def test_returns_binary_image_with_decoded_bytes(self, monkeypatch):
|
|
png = b"\x89PNG\r\n\x1a\nfake-bytes"
|
|
b64 = base64.b64encode(png).decode()
|
|
capture: dict = {}
|
|
_patch_generate(
|
|
monkeypatch, GenerateResponse(response="", done=True, image=b64), capture
|
|
)
|
|
|
|
result = await generate_image(prompt="a red square")
|
|
|
|
assert isinstance(result, BinaryImage)
|
|
assert result.data == png
|
|
assert result.media_type == "image/png"
|
|
assert capture["prompt"] == "a red square"
|
|
assert capture["model"] == envConfig.OTERM_OLLAMA_IMAGE_MODEL
|
|
|
|
async def test_model_argument_overrides_default(self, monkeypatch):
|
|
capture: dict = {}
|
|
_patch_generate(
|
|
monkeypatch,
|
|
GenerateResponse(
|
|
response="", done=True, image=base64.b64encode(b"x").decode()
|
|
),
|
|
capture,
|
|
)
|
|
|
|
await generate_image(prompt="cat", model="x/flux2-klein")
|
|
|
|
assert capture["model"] == "x/flux2-klein"
|
|
|
|
async def test_missing_image_logs_and_raises_model_retry(self, monkeypatch):
|
|
_patch_generate(
|
|
monkeypatch,
|
|
GenerateResponse(response="hi", done=True, image=None),
|
|
{},
|
|
)
|
|
log_lines.clear()
|
|
|
|
with pytest.raises(ModelRetry, match="did not return an image"):
|
|
await generate_image(prompt="cat", model="qwen3")
|
|
|
|
messages = [m for _, m in log_lines]
|
|
assert any("qwen3" in m and "no image" in m.lower() for m in messages)
|
|
|
|
async def test_response_error_is_logged_and_re_raised_as_retry(self, monkeypatch):
|
|
_patch_generate_raises(monkeypatch, ResponseError("model not found", 404))
|
|
log_lines.clear()
|
|
|
|
with pytest.raises(ModelRetry, match="model not found"):
|
|
await generate_image(prompt="cat", model="bogus")
|
|
|
|
messages = [m for _, m in log_lines]
|
|
assert any("bogus" in m and "404" in m for m in messages)
|
|
|
|
async def test_request_error_is_logged_and_re_raised_as_retry(self, monkeypatch):
|
|
_patch_generate_raises(monkeypatch, RequestError("connection refused"))
|
|
log_lines.clear()
|
|
|
|
with pytest.raises(ModelRetry, match="Could not reach Ollama"):
|
|
await generate_image(prompt="cat", model="x/z-image-turbo")
|
|
|
|
messages = [m for _, m in log_lines]
|
|
assert any("connection refused" in m for m in messages)
|