tests/test_image_provider.py

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"}]})