Spaces:
Running
Running
File size: 1,412 Bytes
c8365f5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 | 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"}]})
|