mirror of
https://github.com/NousResearch/hermes-agent.git
synced 2026-05-31 06:51:29 +00:00
9919caff46
* feat(image_gen): add Krea provider plugin (Krea 2 Medium + Large)
New built-in image_gen backend wrapping Krea's Krea 2 foundation
image model family. Auto-discovered like the other image_gen plugins
and appears in 'hermes tools' → Image Generation → Krea.
Krea's API is asynchronous — submit returns a job_id, poll /jobs/{id}
until terminal. The provider hides that behind the synchronous
ImageGenProvider.generate() contract: submit, poll every 2s with
light backoff (max 5s), 3-minute ceiling matching Krea's hosted-tool
timeout. Result URL is materialised to $HERMES_HOME/cache/images/
to avoid CDN-expiry 404s downstream (same fix as xAI #26942).
Models:
- krea-2-medium (default — Krea's 'start here' recommendation)
- krea-2-large
Aspect ratios map landscape→16:9, square→1:1, portrait→9:16.
Resolution: 1K (Krea's only current option).
Kwarg passthrough: seed, creativity (raw/low/medium/high), styles,
image_style_references (capped 10), moodboards (capped 1) — matches
Krea's per-request limits. Unknown kwargs are ignored.
Config knobs (config.yaml):
image_gen.provider: krea
image_gen.krea.model: krea-2-medium | krea-2-large
image_gen.krea.creativity: raw | low | medium | high
Env overrides: KREA_API_KEY (required), KREA_IMAGE_MODEL.
KREA_API_KEY is registered in OPTIONAL_ENV_VARS so 'hermes setup'
prompts for it.
31 new tests; image_gen suite + picker + tools_config: 211/211.
* fix(image_gen/krea): address review feedback
- Update KREA_API_KEY setup URL to the canonical token-creation page
(https://www.krea.ai/app/api/tokens). The previous URL returned 404.
- Fail fast on non-retryable HTTP statuses during poll. The previous
loop retried every HTTPError for the full 180s deadline, so an auth
(401), billing (402), forbidden (403), or not-found (404) response
would make image_generate hang for three minutes. Only retry
transient statuses (408/409/425/429/5xx); surface everything else
immediately.
- Add 5 tests covering fail-fast on 401/403/404 and retry on 429/503.
* fix(krea): point users at the real API token dashboard URL
Three call sites linked users to dashboard pages that don't exist:
- hermes_cli/config.py: https://www.krea.ai/app/api/tokens
- plugins/image_gen/krea/__init__.py get_setup_schema: https://www.krea.ai/api-keys
- plugins/image_gen/krea/__init__.py auth_required error: https://www.krea.ai/api-keys
Per Krea's own docs (https://docs.krea.ai/developers/api-keys-and-billing),
the real dashboard URL is https://www.krea.ai/settings/api-tokens. All three
sites now point there.
625 lines
24 KiB
Python
625 lines
24 KiB
Python
#!/usr/bin/env python3
|
|
"""Tests for Krea image generation provider."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Fixtures
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _fake_api_key(monkeypatch):
|
|
"""Ensure KREA_API_KEY is set for all tests."""
|
|
monkeypatch.setenv("KREA_API_KEY", "test-key-12345")
|
|
|
|
|
|
def _completed_job(url: str = "https://krea.cdn/img.png") -> dict:
|
|
return {
|
|
"job_id": "00000000-0000-0000-0000-000000000abc",
|
|
"status": "completed",
|
|
"created_at": "2026-05-27T00:00:00Z",
|
|
"completed_at": "2026-05-27T00:00:30Z",
|
|
"result": {"urls": [url]},
|
|
}
|
|
|
|
|
|
def _submit_response(job_id: str = "00000000-0000-0000-0000-000000000abc"):
|
|
resp = MagicMock()
|
|
resp.status_code = 200
|
|
resp.raise_for_status = MagicMock()
|
|
resp.json.return_value = {
|
|
"job_id": job_id,
|
|
"status": "queued",
|
|
"created_at": "2026-05-27T00:00:00Z",
|
|
"completed_at": None,
|
|
"result": None,
|
|
}
|
|
return resp
|
|
|
|
|
|
def _poll_response(body: dict):
|
|
resp = MagicMock()
|
|
resp.status_code = 200
|
|
resp.raise_for_status = MagicMock()
|
|
resp.json.return_value = body
|
|
return resp
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Provider class tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestKreaImageGenProvider:
|
|
def test_name(self):
|
|
from plugins.image_gen.krea import KreaImageGenProvider
|
|
|
|
assert KreaImageGenProvider().name == "krea"
|
|
|
|
def test_display_name(self):
|
|
from plugins.image_gen.krea import KreaImageGenProvider
|
|
|
|
assert KreaImageGenProvider().display_name == "Krea"
|
|
|
|
def test_is_available_with_key(self, monkeypatch):
|
|
monkeypatch.setenv("KREA_API_KEY", "sk-test")
|
|
from plugins.image_gen.krea import KreaImageGenProvider
|
|
|
|
assert KreaImageGenProvider().is_available() is True
|
|
|
|
def test_is_available_without_key(self, monkeypatch):
|
|
monkeypatch.delenv("KREA_API_KEY", raising=False)
|
|
from plugins.image_gen.krea import KreaImageGenProvider
|
|
|
|
assert KreaImageGenProvider().is_available() is False
|
|
|
|
def test_list_models(self):
|
|
from plugins.image_gen.krea import KreaImageGenProvider
|
|
|
|
models = KreaImageGenProvider().list_models()
|
|
ids = {m["id"] for m in models}
|
|
assert {"krea-2-medium", "krea-2-large"} <= ids
|
|
# Each entry carries the picker fields the registry expects.
|
|
for m in models:
|
|
assert m["display"]
|
|
assert m["speed"]
|
|
assert m["strengths"]
|
|
assert m["price"]
|
|
|
|
def test_default_model_is_medium(self):
|
|
from plugins.image_gen.krea import KreaImageGenProvider
|
|
|
|
assert KreaImageGenProvider().default_model() == "krea-2-medium"
|
|
|
|
def test_get_setup_schema(self):
|
|
from plugins.image_gen.krea import KreaImageGenProvider
|
|
|
|
schema = KreaImageGenProvider().get_setup_schema()
|
|
assert schema["name"] == "Krea"
|
|
assert schema["badge"] == "paid"
|
|
env_vars = schema["env_vars"]
|
|
assert len(env_vars) == 1
|
|
assert env_vars[0]["key"] == "KREA_API_KEY"
|
|
assert "krea.ai" in env_vars[0]["url"]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Model resolution
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestModelResolution:
|
|
def test_default(self):
|
|
from plugins.image_gen.krea import _resolve_model
|
|
|
|
model_id, meta = _resolve_model()
|
|
assert model_id == "krea-2-medium"
|
|
assert meta["path"] == "medium"
|
|
|
|
def test_env_override_large(self, monkeypatch):
|
|
monkeypatch.setenv("KREA_IMAGE_MODEL", "krea-2-large")
|
|
from plugins.image_gen.krea import _resolve_model
|
|
|
|
model_id, meta = _resolve_model()
|
|
assert model_id == "krea-2-large"
|
|
assert meta["path"] == "large"
|
|
|
|
def test_env_override_unknown_falls_back_to_default(self, monkeypatch):
|
|
monkeypatch.setenv("KREA_IMAGE_MODEL", "krea-2-xxl-fake")
|
|
from plugins.image_gen.krea import _resolve_model
|
|
|
|
model_id, _ = _resolve_model()
|
|
assert model_id == "krea-2-medium"
|
|
|
|
def test_creativity_default(self):
|
|
from plugins.image_gen.krea import _resolve_creativity
|
|
|
|
assert _resolve_creativity(None) == "medium"
|
|
|
|
def test_creativity_valid(self):
|
|
from plugins.image_gen.krea import _resolve_creativity
|
|
|
|
assert _resolve_creativity("HIGH") == "high"
|
|
assert _resolve_creativity(" raw ") == "raw"
|
|
|
|
def test_creativity_invalid(self):
|
|
from plugins.image_gen.krea import _resolve_creativity
|
|
|
|
assert _resolve_creativity("ultra") == "medium"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Generate — main flow
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestGenerate:
|
|
def test_missing_api_key(self, monkeypatch):
|
|
monkeypatch.delenv("KREA_API_KEY", raising=False)
|
|
from plugins.image_gen.krea import KreaImageGenProvider
|
|
|
|
result = KreaImageGenProvider().generate(prompt="test")
|
|
assert result["success"] is False
|
|
assert "KREA_API_KEY" in result["error"]
|
|
assert result["error_type"] == "auth_required"
|
|
|
|
def test_empty_prompt(self):
|
|
from plugins.image_gen.krea import KreaImageGenProvider
|
|
|
|
result = KreaImageGenProvider().generate(prompt=" ")
|
|
assert result["success"] is False
|
|
assert result["error_type"] == "invalid_argument"
|
|
|
|
def test_successful_generation(self):
|
|
"""Happy path: submit → one poll → completed → URL downloaded."""
|
|
from plugins.image_gen.krea import KreaImageGenProvider
|
|
|
|
submit = _submit_response()
|
|
poll = _poll_response(_completed_job("https://krea.cdn/result.png"))
|
|
|
|
with patch("plugins.image_gen.krea.requests.post", return_value=submit) as mock_post, \
|
|
patch("plugins.image_gen.krea.requests.get", return_value=poll) as mock_get, \
|
|
patch(
|
|
"plugins.image_gen.krea.save_url_image",
|
|
return_value=Path("/tmp/krea_krea-2-medium_test.png"),
|
|
) as mock_save, \
|
|
patch("plugins.image_gen.krea.time.sleep"): # skip real waits
|
|
result = KreaImageGenProvider().generate(prompt="A cinematic lamp")
|
|
|
|
assert result["success"] is True
|
|
assert result["image"] == "/tmp/krea_krea-2-medium_test.png"
|
|
assert result["provider"] == "krea"
|
|
assert result["model"] == "krea-2-medium"
|
|
assert result["aspect_ratio"] == "landscape"
|
|
assert result["job_id"] == "00000000-0000-0000-0000-000000000abc"
|
|
assert result["resolution"] == "1K"
|
|
assert result["creativity"] == "medium"
|
|
# Submit hit the medium endpoint
|
|
post_url = mock_post.call_args[0][0]
|
|
assert post_url.endswith("/generate/image/krea/krea-2/medium")
|
|
# Poll hit /jobs/{job_id}
|
|
poll_url = mock_get.call_args[0][0]
|
|
assert "/jobs/00000000-0000-0000-0000-000000000abc" in poll_url
|
|
# URL was materialised once
|
|
mock_save.assert_called_once()
|
|
|
|
def test_large_model_routes_to_large_endpoint(self, monkeypatch):
|
|
monkeypatch.setenv("KREA_IMAGE_MODEL", "krea-2-large")
|
|
from plugins.image_gen.krea import KreaImageGenProvider
|
|
|
|
submit = _submit_response()
|
|
poll = _poll_response(_completed_job())
|
|
|
|
with patch("plugins.image_gen.krea.requests.post", return_value=submit) as mock_post, \
|
|
patch("plugins.image_gen.krea.requests.get", return_value=poll), \
|
|
patch(
|
|
"plugins.image_gen.krea.save_url_image",
|
|
return_value=Path("/tmp/x.png"),
|
|
), \
|
|
patch("plugins.image_gen.krea.time.sleep"):
|
|
KreaImageGenProvider().generate(prompt="test")
|
|
|
|
post_url = mock_post.call_args[0][0]
|
|
assert post_url.endswith("/generate/image/krea/krea-2/large")
|
|
|
|
def test_aspect_ratio_mapping(self):
|
|
"""Hermes 'square' must map to Krea '1:1' in the wire payload."""
|
|
from plugins.image_gen.krea import KreaImageGenProvider
|
|
|
|
submit = _submit_response()
|
|
poll = _poll_response(_completed_job())
|
|
|
|
with patch("plugins.image_gen.krea.requests.post", return_value=submit) as mock_post, \
|
|
patch("plugins.image_gen.krea.requests.get", return_value=poll), \
|
|
patch(
|
|
"plugins.image_gen.krea.save_url_image",
|
|
return_value=Path("/tmp/x.png"),
|
|
), \
|
|
patch("plugins.image_gen.krea.time.sleep"):
|
|
KreaImageGenProvider().generate(prompt="test", aspect_ratio="square")
|
|
|
|
payload = mock_post.call_args.kwargs["json"]
|
|
assert payload["aspect_ratio"] == "1:1"
|
|
assert payload["resolution"] == "1K"
|
|
|
|
def test_auth_header(self):
|
|
from plugins.image_gen.krea import KreaImageGenProvider
|
|
|
|
submit = _submit_response()
|
|
poll = _poll_response(_completed_job())
|
|
|
|
with patch("plugins.image_gen.krea.requests.post", return_value=submit) as mock_post, \
|
|
patch("plugins.image_gen.krea.requests.get", return_value=poll), \
|
|
patch(
|
|
"plugins.image_gen.krea.save_url_image",
|
|
return_value=Path("/tmp/x.png"),
|
|
), \
|
|
patch("plugins.image_gen.krea.time.sleep"):
|
|
KreaImageGenProvider().generate(prompt="test")
|
|
|
|
headers = mock_post.call_args.kwargs["headers"]
|
|
assert headers["Authorization"] == "Bearer test-key-12345"
|
|
assert headers["Content-Type"] == "application/json"
|
|
|
|
def test_passthrough_seed_styles_moodboards(self):
|
|
from plugins.image_gen.krea import KreaImageGenProvider
|
|
|
|
submit = _submit_response()
|
|
poll = _poll_response(_completed_job())
|
|
|
|
with patch("plugins.image_gen.krea.requests.post", return_value=submit) as mock_post, \
|
|
patch("plugins.image_gen.krea.requests.get", return_value=poll), \
|
|
patch(
|
|
"plugins.image_gen.krea.save_url_image",
|
|
return_value=Path("/tmp/x.png"),
|
|
), \
|
|
patch("plugins.image_gen.krea.time.sleep"):
|
|
KreaImageGenProvider().generate(
|
|
prompt="test",
|
|
seed=42,
|
|
styles=[{"id": "lora-1", "strength": 0.7}],
|
|
moodboards=[{"url": "https://x.com/mood.png"}, {"url": "https://x.com/mood2.png"}],
|
|
image_style_references=[{"url": f"https://x.com/{i}.png"} for i in range(15)],
|
|
creativity="high",
|
|
)
|
|
|
|
payload = mock_post.call_args.kwargs["json"]
|
|
assert payload["seed"] == 42
|
|
assert payload["styles"] == [{"id": "lora-1", "strength": 0.7}]
|
|
assert len(payload["moodboards"]) == 1 # capped at 1
|
|
assert len(payload["image_style_references"]) == 10 # capped at 10
|
|
assert payload["creativity"] == "high"
|
|
|
|
def test_unknown_kwargs_ignored(self):
|
|
"""Forward-compat: unknown kwargs must not break generate()."""
|
|
from plugins.image_gen.krea import KreaImageGenProvider
|
|
|
|
submit = _submit_response()
|
|
poll = _poll_response(_completed_job())
|
|
|
|
with patch("plugins.image_gen.krea.requests.post", return_value=submit), \
|
|
patch("plugins.image_gen.krea.requests.get", return_value=poll), \
|
|
patch(
|
|
"plugins.image_gen.krea.save_url_image",
|
|
return_value=Path("/tmp/x.png"),
|
|
), \
|
|
patch("plugins.image_gen.krea.time.sleep"):
|
|
result = KreaImageGenProvider().generate(
|
|
prompt="test",
|
|
fictional_param="should be ignored",
|
|
num_images=4,
|
|
)
|
|
|
|
assert result["success"] is True
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Generate — error paths
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestGenerateErrors:
|
|
def test_submit_http_error(self):
|
|
import requests as req_lib
|
|
from plugins.image_gen.krea import KreaImageGenProvider
|
|
|
|
resp = req_lib.Response()
|
|
resp.status_code = 401
|
|
resp._content = b'{"error": {"message": "Invalid API key"}}'
|
|
resp.headers["Content-Type"] = "application/json"
|
|
resp.raise_for_status = MagicMock(
|
|
side_effect=req_lib.HTTPError(response=resp)
|
|
)
|
|
|
|
with patch("plugins.image_gen.krea.requests.post", return_value=resp):
|
|
result = KreaImageGenProvider().generate(prompt="test")
|
|
|
|
assert result["success"] is False
|
|
assert result["error_type"] == "api_error"
|
|
assert "401" in result["error"]
|
|
assert "Invalid API key" in result["error"]
|
|
|
|
def test_submit_timeout(self):
|
|
import requests as req_lib
|
|
from plugins.image_gen.krea import KreaImageGenProvider
|
|
|
|
with patch(
|
|
"plugins.image_gen.krea.requests.post", side_effect=req_lib.Timeout()
|
|
):
|
|
result = KreaImageGenProvider().generate(prompt="test")
|
|
|
|
assert result["success"] is False
|
|
assert result["error_type"] == "timeout"
|
|
|
|
def test_submit_connection_error(self):
|
|
import requests as req_lib
|
|
from plugins.image_gen.krea import KreaImageGenProvider
|
|
|
|
with patch(
|
|
"plugins.image_gen.krea.requests.post",
|
|
side_effect=req_lib.ConnectionError("dns nope"),
|
|
):
|
|
result = KreaImageGenProvider().generate(prompt="test")
|
|
|
|
assert result["success"] is False
|
|
assert result["error_type"] == "connection_error"
|
|
|
|
def test_submit_missing_job_id(self):
|
|
from plugins.image_gen.krea import KreaImageGenProvider
|
|
|
|
bad_submit = MagicMock()
|
|
bad_submit.status_code = 200
|
|
bad_submit.raise_for_status = MagicMock()
|
|
bad_submit.json.return_value = {"status": "queued"}
|
|
|
|
with patch("plugins.image_gen.krea.requests.post", return_value=bad_submit):
|
|
result = KreaImageGenProvider().generate(prompt="test")
|
|
|
|
assert result["success"] is False
|
|
assert result["error_type"] == "invalid_response"
|
|
assert "job_id" in result["error"]
|
|
|
|
def test_job_failed(self):
|
|
from plugins.image_gen.krea import KreaImageGenProvider
|
|
|
|
failed = {
|
|
"job_id": "abc",
|
|
"status": "failed",
|
|
"completed_at": "2026-05-27T00:01:00Z",
|
|
"result": {"error": "NSFW content"},
|
|
}
|
|
|
|
submit = _submit_response()
|
|
with patch("plugins.image_gen.krea.requests.post", return_value=submit), \
|
|
patch(
|
|
"plugins.image_gen.krea.requests.get",
|
|
return_value=_poll_response(failed),
|
|
), \
|
|
patch("plugins.image_gen.krea.time.sleep"):
|
|
result = KreaImageGenProvider().generate(prompt="test")
|
|
|
|
assert result["success"] is False
|
|
assert result["error_type"] == "api_error"
|
|
assert "NSFW" in result["error"]
|
|
|
|
def test_job_cancelled(self):
|
|
from plugins.image_gen.krea import KreaImageGenProvider
|
|
|
|
cancelled = {
|
|
"job_id": "abc",
|
|
"status": "cancelled",
|
|
"completed_at": "2026-05-27T00:01:00Z",
|
|
"result": {},
|
|
}
|
|
|
|
with patch("plugins.image_gen.krea.requests.post", return_value=_submit_response()), \
|
|
patch(
|
|
"plugins.image_gen.krea.requests.get",
|
|
return_value=_poll_response(cancelled),
|
|
), \
|
|
patch("plugins.image_gen.krea.time.sleep"):
|
|
result = KreaImageGenProvider().generate(prompt="test")
|
|
|
|
assert result["success"] is False
|
|
assert result["error_type"] == "cancelled"
|
|
|
|
def test_completed_but_missing_urls(self):
|
|
from plugins.image_gen.krea import KreaImageGenProvider
|
|
|
|
completed_empty = {
|
|
"job_id": "abc",
|
|
"status": "completed",
|
|
"completed_at": "2026-05-27T00:01:00Z",
|
|
"result": {"urls": []},
|
|
}
|
|
|
|
with patch("plugins.image_gen.krea.requests.post", return_value=_submit_response()), \
|
|
patch(
|
|
"plugins.image_gen.krea.requests.get",
|
|
return_value=_poll_response(completed_empty),
|
|
), \
|
|
patch("plugins.image_gen.krea.time.sleep"):
|
|
result = KreaImageGenProvider().generate(prompt="test")
|
|
|
|
assert result["success"] is False
|
|
assert result["error_type"] == "empty_response"
|
|
|
|
def test_url_download_failure_falls_back_to_bare_url(self):
|
|
"""Mirror of xAI behaviour — if local cache fails, return the URL."""
|
|
import requests as req_lib
|
|
from plugins.image_gen.krea import KreaImageGenProvider
|
|
|
|
url = "https://krea.cdn/expired-soon.png"
|
|
submit = _submit_response()
|
|
poll = _poll_response(_completed_job(url))
|
|
|
|
with patch("plugins.image_gen.krea.requests.post", return_value=submit), \
|
|
patch("plugins.image_gen.krea.requests.get", return_value=poll), \
|
|
patch(
|
|
"plugins.image_gen.krea.save_url_image",
|
|
side_effect=req_lib.HTTPError("404"),
|
|
), \
|
|
patch("plugins.image_gen.krea.time.sleep"):
|
|
result = KreaImageGenProvider().generate(prompt="test")
|
|
|
|
assert result["success"] is True
|
|
assert result["image"] == url
|
|
|
|
def test_polling_picks_up_completed_at_with_unknown_status(self):
|
|
"""``completed_at`` set + unrecognised pending status → still terminal."""
|
|
from plugins.image_gen.krea import KreaImageGenProvider
|
|
|
|
# Use a status value that is NOT in our terminal set ("intermediate-complete")
|
|
# but with completed_at populated — Krea's spec says completed_at is the
|
|
# canonical terminal marker.
|
|
oddball = {
|
|
"job_id": "abc",
|
|
"status": "intermediate-complete",
|
|
"completed_at": "2026-05-27T00:01:00Z",
|
|
"result": {"urls": ["https://krea.cdn/done.png"]},
|
|
}
|
|
|
|
with patch("plugins.image_gen.krea.requests.post", return_value=_submit_response()), \
|
|
patch(
|
|
"plugins.image_gen.krea.requests.get",
|
|
return_value=_poll_response(oddball),
|
|
), \
|
|
patch(
|
|
"plugins.image_gen.krea.save_url_image",
|
|
return_value=Path("/tmp/x.png"),
|
|
), \
|
|
patch("plugins.image_gen.krea.time.sleep"):
|
|
result = KreaImageGenProvider().generate(prompt="test")
|
|
|
|
assert result["success"] is True
|
|
|
|
|
|
class TestPollRetryPolicy:
|
|
"""Polling fail-fast on permanent 4xx, retry on transient 5xx/429."""
|
|
|
|
def _http_error_response(self, status: int):
|
|
import requests as req_lib
|
|
|
|
resp = req_lib.Response()
|
|
resp.status_code = status
|
|
resp._content = b'{"error": "boom"}'
|
|
resp.headers["Content-Type"] = "application/json"
|
|
resp.raise_for_status = MagicMock(
|
|
side_effect=req_lib.HTTPError(response=resp)
|
|
)
|
|
return resp
|
|
|
|
def test_poll_fails_fast_on_401(self):
|
|
"""Auth failure mid-poll should not wait the 180s deadline."""
|
|
from plugins.image_gen.krea import KreaImageGenProvider
|
|
|
|
bad_poll = self._http_error_response(401)
|
|
|
|
with patch("plugins.image_gen.krea.requests.post", return_value=_submit_response()), \
|
|
patch("plugins.image_gen.krea.requests.get", return_value=bad_poll) as mock_get, \
|
|
patch("plugins.image_gen.krea.time.sleep"):
|
|
result = KreaImageGenProvider().generate(prompt="test")
|
|
|
|
assert result["success"] is False
|
|
assert result["error_type"] == "api_error"
|
|
assert "401" in result["error"]
|
|
# One call — no retry on permanent auth failure.
|
|
assert mock_get.call_count == 1
|
|
|
|
def test_poll_fails_fast_on_404(self):
|
|
"""Missing job (404) should surface immediately, not retry for 180s."""
|
|
from plugins.image_gen.krea import KreaImageGenProvider
|
|
|
|
bad_poll = self._http_error_response(404)
|
|
|
|
with patch("plugins.image_gen.krea.requests.post", return_value=_submit_response()), \
|
|
patch("plugins.image_gen.krea.requests.get", return_value=bad_poll) as mock_get, \
|
|
patch("plugins.image_gen.krea.time.sleep"):
|
|
result = KreaImageGenProvider().generate(prompt="test")
|
|
|
|
assert result["success"] is False
|
|
assert result["error_type"] == "api_error"
|
|
assert "404" in result["error"]
|
|
assert mock_get.call_count == 1
|
|
|
|
def test_poll_fails_fast_on_403(self):
|
|
"""Billing/permission failure (403) should not retry."""
|
|
from plugins.image_gen.krea import KreaImageGenProvider
|
|
|
|
bad_poll = self._http_error_response(403)
|
|
|
|
with patch("plugins.image_gen.krea.requests.post", return_value=_submit_response()), \
|
|
patch("plugins.image_gen.krea.requests.get", return_value=bad_poll) as mock_get, \
|
|
patch("plugins.image_gen.krea.time.sleep"):
|
|
result = KreaImageGenProvider().generate(prompt="test")
|
|
|
|
assert result["success"] is False
|
|
assert mock_get.call_count == 1
|
|
|
|
def test_poll_retries_on_503_then_succeeds(self):
|
|
"""Transient 5xx should retry and eventually surface a completion."""
|
|
from plugins.image_gen.krea import KreaImageGenProvider
|
|
|
|
flaky = self._http_error_response(503)
|
|
good = _poll_response(_completed_job("https://krea.cdn/ok.png"))
|
|
|
|
with patch("plugins.image_gen.krea.requests.post", return_value=_submit_response()), \
|
|
patch(
|
|
"plugins.image_gen.krea.requests.get",
|
|
side_effect=[flaky, flaky, good],
|
|
) as mock_get, \
|
|
patch(
|
|
"plugins.image_gen.krea.save_url_image",
|
|
return_value=Path("/tmp/x.png"),
|
|
), \
|
|
patch("plugins.image_gen.krea.time.sleep"):
|
|
result = KreaImageGenProvider().generate(prompt="test")
|
|
|
|
assert result["success"] is True
|
|
assert mock_get.call_count == 3
|
|
|
|
def test_poll_retries_on_429(self):
|
|
"""Rate-limit (429) is in the retryable set."""
|
|
from plugins.image_gen.krea import KreaImageGenProvider
|
|
|
|
rate_limited = self._http_error_response(429)
|
|
good = _poll_response(_completed_job("https://krea.cdn/ok.png"))
|
|
|
|
with patch("plugins.image_gen.krea.requests.post", return_value=_submit_response()), \
|
|
patch(
|
|
"plugins.image_gen.krea.requests.get",
|
|
side_effect=[rate_limited, good],
|
|
) as mock_get, \
|
|
patch(
|
|
"plugins.image_gen.krea.save_url_image",
|
|
return_value=Path("/tmp/x.png"),
|
|
), \
|
|
patch("plugins.image_gen.krea.time.sleep"):
|
|
result = KreaImageGenProvider().generate(prompt="test")
|
|
|
|
assert result["success"] is True
|
|
assert mock_get.call_count == 2
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Registration
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestRegistration:
|
|
def test_register(self):
|
|
from plugins.image_gen.krea import KreaImageGenProvider, register
|
|
|
|
mock_ctx = MagicMock()
|
|
register(mock_ctx)
|
|
mock_ctx.register_image_gen_provider.assert_called_once()
|
|
provider = mock_ctx.register_image_gen_provider.call_args[0][0]
|
|
assert isinstance(provider, KreaImageGenProvider)
|
|
assert provider.name == "krea"
|