44 lines
1.5 KiB
Python
44 lines
1.5 KiB
Python
import base64
|
|
|
|
import pytest
|
|
from fastapi import HTTPException
|
|
|
|
from openocr_markdown.openai_schemas import ChatMessage, ChatMessageContentImageUrl, ChatMessageContentPart
|
|
from openocr_markdown.openai_utils import DATA_URL_PATTERN, extract_image_urls, materialize_image_url
|
|
|
|
|
|
def test_extract_image_urls_from_multimodal_message() -> None:
|
|
messages = [
|
|
ChatMessage(
|
|
role="user",
|
|
content=[
|
|
ChatMessageContentPart(type="text", text="Convert to markdown"),
|
|
ChatMessageContentPart(
|
|
type="image_url",
|
|
image_url=ChatMessageContentImageUrl(url="https://example.com/page.png"),
|
|
),
|
|
],
|
|
)
|
|
]
|
|
assert extract_image_urls(messages) == ["https://example.com/page.png"]
|
|
|
|
|
|
def test_materialize_data_url(tmp_path) -> None:
|
|
png_bytes = base64.b64decode(
|
|
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg=="
|
|
)
|
|
data_url = "data:image/png;base64," + base64.b64encode(png_bytes).decode()
|
|
path = materialize_image_url(data_url, tmp_path)
|
|
assert path.exists()
|
|
assert path.suffix == ".png"
|
|
|
|
|
|
def test_data_url_pattern_rejects_invalid_input() -> None:
|
|
assert DATA_URL_PATTERN.match("not-a-data-url") is None
|
|
|
|
|
|
def test_materialize_image_url_rejects_unsupported_scheme(tmp_path) -> None:
|
|
with pytest.raises(HTTPException) as exc:
|
|
materialize_image_url("ftp://example.com/a.png", tmp_path)
|
|
assert exc.value.status_code == 400
|