from __future__ import annotations import pytest from api.image_provider import ImageProviderError, _read_remote_image, is_pollinations_remote_url def test_accepts_only_https_pollinations_urls() -> None: assert is_pollinations_remote_url("https://image.pollinations.ai/prompt/cat?seed=1") assert is_pollinations_remote_url("https://gen.pollinations.ai/image/cat") assert not is_pollinations_remote_url("http://image.pollinations.ai/prompt/cat") assert not is_pollinations_remote_url("https://example.com/image.jpg") assert not is_pollinations_remote_url("data:image/jpeg;base64,AA==") assert not is_pollinations_remote_url("https://image.pollinations.ai.evil.example/image") def test_parses_safe_remote_result() -> None: image = _read_remote_image({ "data": [{ "url": "https://image.pollinations.ai/prompt/a-safe-cat?seed=7", "media_type": "image/png", "revised_prompt": "a safe cat", }] }) assert image.url.startswith("https://image.pollinations.ai/") assert image.mime_type == "image/png" assert image.revised_prompt == "a safe cat" def test_rejects_missing_or_untrusted_remote_result() -> None: with pytest.raises(ImageProviderError): _read_remote_image({"data": []}) with pytest.raises(ImageProviderError): _read_remote_image({"data": [{"url": "https://example.com/image.jpg"}]})