wehub-resource-sync
218 lines
7.6 KiB
Python
218 lines
7.6 KiB
Python
from __future__ import annotations
|
|
|
|
import io
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
from PIL import Image
|
|
|
|
from opensquilla.artifacts import (
|
|
ARTIFACT_THUMBNAIL_NAME,
|
|
ArtifactStore,
|
|
artifact_payload,
|
|
)
|
|
|
|
|
|
class _FakeSessionManager:
|
|
def __init__(self, session_id: str) -> None:
|
|
self.session_id = session_id
|
|
|
|
async def get_session(self, session_key: str) -> object | None:
|
|
if session_key == "agent:main:webchat:ok":
|
|
return SimpleNamespace(session_id=self.session_id)
|
|
return None
|
|
|
|
|
|
def _app(tmp_path: Path):
|
|
pytest.importorskip("starlette.testclient")
|
|
from starlette.applications import Starlette
|
|
|
|
from opensquilla.gateway.artifacts import register_artifact_routes
|
|
from opensquilla.gateway.config import AttachmentsConfig, AuthConfig, GatewayConfig
|
|
from opensquilla.gateway.middleware import AuthMiddleware
|
|
|
|
config = GatewayConfig(
|
|
auth=AuthConfig(mode="token", token="secret"),
|
|
attachments=AttachmentsConfig(media_root=str(tmp_path)),
|
|
)
|
|
app = Starlette(debug=False)
|
|
register_artifact_routes(
|
|
app,
|
|
config=config,
|
|
session_manager=_FakeSessionManager("session-1"),
|
|
)
|
|
app.add_middleware(AuthMiddleware, config=config)
|
|
return app
|
|
|
|
|
|
def _png_bytes(width: int = 1024, height: int = 768) -> bytes:
|
|
image = Image.new("RGB", (width, height), color=(180, 40, 90))
|
|
# Add a little structure so the encoder produces a non-trivial-sized image.
|
|
for x in range(0, width, 8):
|
|
for y in range(0, height, 8):
|
|
image.putpixel((x, y), ((x * 3) % 256, (y * 5) % 256, (x + y) % 256))
|
|
out = io.BytesIO()
|
|
image.save(out, format="PNG")
|
|
return out.getvalue()
|
|
|
|
|
|
def _publish_image(tmp_path: Path):
|
|
payload = _png_bytes()
|
|
ref = ArtifactStore(tmp_path).publish_bytes(
|
|
payload,
|
|
session_id="session-1",
|
|
session_key="agent:main:webchat:ok",
|
|
name="render.png",
|
|
mime="image/png",
|
|
source="image_generate",
|
|
)
|
|
return ref, payload
|
|
|
|
|
|
def test_publishing_image_creates_smaller_webp_thumbnail(tmp_path: Path) -> None:
|
|
ref, payload = _publish_image(tmp_path)
|
|
|
|
assert ref.has_thumbnail is True
|
|
thumb_path = ArtifactStore(tmp_path).thumbnail_path_for(ref)
|
|
assert thumb_path.name == ARTIFACT_THUMBNAIL_NAME
|
|
assert thumb_path.exists()
|
|
|
|
thumb_bytes = thumb_path.read_bytes()
|
|
assert thumb_bytes[:4] == b"RIFF" # webp container magic
|
|
assert b"WEBP" in thumb_bytes[:16]
|
|
assert len(thumb_bytes) < len(payload)
|
|
|
|
# The thumbnail must fit within the 512x512 bounding box, preserving aspect.
|
|
with Image.open(io.BytesIO(thumb_bytes)) as decoded:
|
|
assert decoded.width <= 512
|
|
assert decoded.height <= 512
|
|
assert max(decoded.width, decoded.height) == 512
|
|
|
|
|
|
def test_artifact_payload_exposes_thumbnail_url_only_when_present(tmp_path: Path) -> None:
|
|
ref, _payload = _publish_image(tmp_path)
|
|
|
|
payload = artifact_payload(ref)
|
|
assert payload["thumbnail_url"] == f"{payload['download_url']}?variant=thumb"
|
|
assert payload["thumbnail_url"] == f"/api/v1/artifacts/{ref.id}?variant=thumb"
|
|
|
|
|
|
def test_thumbnail_served_at_variant_thumb_and_is_smaller(tmp_path: Path) -> None:
|
|
pytest.importorskip("starlette.testclient")
|
|
from starlette.testclient import TestClient
|
|
|
|
ref, _payload = _publish_image(tmp_path)
|
|
|
|
with TestClient(_app(tmp_path)) as client:
|
|
full = client.get(
|
|
f"/api/v1/artifacts/{ref.id}?sessionKey=agent:main:webchat:ok",
|
|
headers={"Authorization": "Bearer secret"},
|
|
)
|
|
thumb = client.get(
|
|
f"/api/v1/artifacts/{ref.id}?variant=thumb&sessionKey=agent:main:webchat:ok",
|
|
headers={"Authorization": "Bearer secret"},
|
|
)
|
|
thumb_head = client.head(
|
|
f"/api/v1/artifacts/{ref.id}?variant=thumb&sessionKey=agent:main:webchat:ok",
|
|
headers={"Authorization": "Bearer secret"},
|
|
)
|
|
|
|
assert full.status_code == 200
|
|
assert full.headers["content-type"].startswith("image/png")
|
|
|
|
assert thumb.status_code == 200
|
|
assert thumb.headers["content-type"].startswith("image/webp")
|
|
assert thumb.content[:4] == b"RIFF"
|
|
assert int(thumb.headers["content-length"]) < int(full.headers["content-length"])
|
|
|
|
assert thumb_head.status_code == 200
|
|
assert thumb_head.headers["content-type"].startswith("image/webp")
|
|
|
|
|
|
def test_non_image_artifact_has_no_thumbnail_and_falls_back(tmp_path: Path) -> None:
|
|
pytest.importorskip("starlette.testclient")
|
|
from starlette.testclient import TestClient
|
|
|
|
ref = ArtifactStore(tmp_path).publish_bytes(
|
|
b"plain text deliverable contents",
|
|
session_id="session-1",
|
|
session_key="agent:main:webchat:ok",
|
|
name="notes.txt",
|
|
mime="text/plain",
|
|
source="publish_artifact",
|
|
)
|
|
|
|
assert ref.has_thumbnail is False
|
|
assert not ArtifactStore(tmp_path).thumbnail_path_for(ref).exists()
|
|
assert "thumbnail_url" not in artifact_payload(ref)
|
|
|
|
with TestClient(_app(tmp_path)) as client:
|
|
thumb = client.get(
|
|
f"/api/v1/artifacts/{ref.id}?variant=thumb&sessionKey=agent:main:webchat:ok",
|
|
headers={"Authorization": "Bearer secret"},
|
|
)
|
|
|
|
# No thumbnail sidecar: the request transparently falls back to the full file.
|
|
assert thumb.status_code == 200
|
|
assert thumb.headers["content-type"].startswith("text/plain")
|
|
assert thumb.content == b"plain text deliverable contents"
|
|
|
|
|
|
def test_corrupt_image_still_publishes_without_thumbnail(tmp_path: Path) -> None:
|
|
pytest.importorskip("starlette.testclient")
|
|
from starlette.testclient import TestClient
|
|
|
|
# Claims to be a PNG but the bytes are not decodable as an image.
|
|
ref = ArtifactStore(tmp_path).publish_bytes(
|
|
b"\x89PNG\r\n\x1a\n not a real image body",
|
|
session_id="session-1",
|
|
session_key="agent:main:webchat:ok",
|
|
name="broken.png",
|
|
mime="image/png",
|
|
source="image_generate",
|
|
)
|
|
|
|
assert ref.has_thumbnail is False
|
|
assert not ArtifactStore(tmp_path).thumbnail_path_for(ref).exists()
|
|
assert "thumbnail_url" not in artifact_payload(ref)
|
|
|
|
with TestClient(_app(tmp_path)) as client:
|
|
# The artifact is fully published and downloadable despite the bad image.
|
|
full = client.get(
|
|
f"/api/v1/artifacts/{ref.id}?sessionKey=agent:main:webchat:ok",
|
|
headers={"Authorization": "Bearer secret"},
|
|
)
|
|
thumb = client.get(
|
|
f"/api/v1/artifacts/{ref.id}?variant=thumb&sessionKey=agent:main:webchat:ok",
|
|
headers={"Authorization": "Bearer secret"},
|
|
)
|
|
|
|
assert full.status_code == 200
|
|
assert thumb.status_code == 200
|
|
# Falls back to the full bytes since no thumbnail exists.
|
|
assert thumb.headers["content-type"].startswith("image/png")
|
|
|
|
|
|
def test_old_artifact_without_thumbnail_field_round_trips(tmp_path: Path) -> None:
|
|
# Simulate a meta.json written before the thumbnail feature existed.
|
|
from opensquilla.artifacts import ArtifactRef
|
|
|
|
legacy_meta = {
|
|
"id": "art-legacyxxxxxxxxxxxxxxxxxxx",
|
|
"sha256": "a" * 64,
|
|
"name": "old.png",
|
|
"mime": "image/png",
|
|
"size": 123,
|
|
"session_id": "session-1",
|
|
"session_key": "agent:main:webchat:ok",
|
|
"source": "image_generate",
|
|
"created_at": "2025-01-01T00:00:00Z",
|
|
"download_url": "/api/v1/artifacts/art-legacyxxxxxxxxxxxxxxxxxxx",
|
|
"kind": "artifact_ref",
|
|
"store": "artifacts",
|
|
}
|
|
ref = ArtifactRef.from_dict(legacy_meta)
|
|
assert ref.has_thumbnail is False
|
|
assert "thumbnail_url" not in artifact_payload(ref)
|