| 1 | import base64 |
| 2 | import tempfile |
| 3 | import unittest |
| 4 | from io import BytesIO |
| 5 | from pathlib import Path |
| 6 | from unittest.mock import AsyncMock, patch |
| 7 | |
| 8 | from PIL import Image |
| 9 | |
| 10 | from agent_runtime.vimax_adapters import _build_image_generator |
| 11 | from tools.image_generator_openrouter_api import ( |
| 12 | ImageGeneratorOpenRouterAPI, |
| 13 | OpenRouterImageAPIError, |
| 14 | _is_retryable_image_error, |
| 15 | ) |
| 16 | |
| 17 | |
| 18 | def _encoded_png(size: tuple[int, int] = (16, 9)) -> str: |
| 19 | buffer = BytesIO() |
| 20 | Image.new("RGB", size, "blue").save(buffer, format="PNG") |
| 21 | return base64.b64encode(buffer.getvalue()).decode("ascii") |
| 22 | |
| 23 | |
| 24 | class OpenRouterImageGeneratorTests(unittest.IsolatedAsyncioTestCase): |
| 25 | async def test_generates_image_with_dedicated_images_api(self): |
| 26 | captured = {} |
| 27 | |
| 28 | async def fake_post(url, *, headers, payload, timeout): |
| 29 | captured.update(url=url, headers=headers, payload=payload, timeout=timeout) |
| 30 | return 200, {"data": [{"b64_json": _encoded_png(), "media_type": "image/png"}]} |
| 31 | |
| 32 | progress = [] |
| 33 | generator = ImageGeneratorOpenRouterAPI(api_key="secret", model="openai/gpt-image-2") |
| 34 | with patch("tools.image_generator_openrouter_api._post_json", fake_post): |
| 35 | result = await generator.generate_single_image( |
| 36 | "a cinematic beach", |
| 37 | aspect_ratio="16:9", |
| 38 | progress=lambda stage, message, metadata: progress.append((stage, message, metadata)), |
| 39 | ) |
| 40 | |
| 41 | self.assertEqual(captured["url"], "https://openrouter.ai/api/v1/images") |
| 42 | self.assertEqual(captured["headers"]["Authorization"], "Bearer secret") |
| 43 | self.assertEqual(captured["payload"]["model"], "openai/gpt-image-2") |
| 44 | self.assertNotIn("aspect_ratio", captured["payload"]) |
| 45 | self.assertIn("landscape image", captured["payload"]["prompt"]) |
| 46 | self.assertEqual(result.data.size, (16, 9)) |
| 47 | self.assertEqual(result.ext, "png") |
| 48 | self.assertEqual([item[0] for item in progress], ["image_generation", "image_completed"]) |
| 49 | |
| 50 | async def test_reference_images_use_data_urls(self): |
| 51 | with tempfile.TemporaryDirectory() as tmp: |
| 52 | reference_path = Path(tmp) / "reference.png" |
| 53 | Image.new("RGB", (16, 9), "red").save(reference_path) |
| 54 | post = AsyncMock(return_value=(200, {"data": [{"b64_json": _encoded_png()}]})) |
| 55 | generator = ImageGeneratorOpenRouterAPI(api_key="secret") |
| 56 | with patch("tools.image_generator_openrouter_api._post_json", post): |
| 57 | await generator.generate_single_image("edit this", [str(reference_path)]) |
| 58 | |
| 59 | payload = post.await_args.kwargs["payload"] |
| 60 | reference_url = payload["input_references"][0]["image_url"]["url"] |
| 61 | self.assertTrue(reference_url.startswith("data:image/png;base64,")) |
| 62 | |
| 63 | async def test_non_retryable_client_error_is_not_repeated(self): |
| 64 | post = AsyncMock(return_value=(400, {"error": {"message": "bad request"}})) |
| 65 | generator = ImageGeneratorOpenRouterAPI(api_key="secret") |
| 66 | with patch("tools.image_generator_openrouter_api._post_json", post): |
| 67 | with self.assertRaises(OpenRouterImageAPIError): |
| 68 | await generator.generate_single_image("bad request") |
| 69 | self.assertEqual(post.await_count, 1) |
| 70 | |
| 71 | def test_retry_policy_is_bounded_to_transient_errors_and_portrait_outputs(self): |
| 72 | self.assertTrue(_is_retryable_image_error(OpenRouterImageAPIError(429, {}))) |
| 73 | self.assertTrue(_is_retryable_image_error(OpenRouterImageAPIError(500, {}))) |
| 74 | self.assertTrue(_is_retryable_image_error(ValueError("Generated image is portrait-oriented (9x16); retrying for a landscape frame"))) |
| 75 | self.assertFalse(_is_retryable_image_error(OpenRouterImageAPIError(401, {}))) |
| 76 | self.assertFalse(_is_retryable_image_error(ValueError("invalid reference image"))) |
| 77 | |
| 78 | def test_agent_factory_selects_openrouter_from_image_base_url(self): |
| 79 | with patch("agent_runtime.vimax_adapters.image_api_key", return_value="secret"), \ |
| 80 | patch("agent_runtime.vimax_adapters.image_model", return_value="openai/gpt-image-2"), \ |
| 81 | patch("agent_runtime.vimax_adapters.image_base_url", return_value="https://openrouter.ai/api/v1"): |
| 82 | generator = _build_image_generator() |
| 83 | self.assertIsInstance(generator, ImageGeneratorOpenRouterAPI) |
| 84 | self.assertEqual(generator.model, "openai/gpt-image-2") |
| 85 | |
| 86 | |
| 87 | if __name__ == "__main__": |
| 88 | unittest.main() |
| 89 |