script-build-and-framework-tmp/agent/tests/test_content_collage.py

from PIL import Image
import pytest

from agent.tools.builtin.content import collage
from agent.tools.builtin.content.platforms import aigc_channel, x, youtube


@pytest.mark.asyncio
async def test_collage_filters_failed_images_and_keeps_matching_labels(monkeypatch):
    image = Image.new("RGB", (4, 4), "white")
    captured = {}

    async def fake_load(urls):
        return [(urls[0], image), (urls[1], None)]

    def fake_grid(*, images, labels):
        captured["images"] = images
        captured["labels"] = labels
        return image

    async def fake_upload(image_bytes, filename):
        captured["filename"] = filename
        return "https://cdn.example/collage.png"

    monkeypatch.setattr(collage, "load_images", fake_load)
    monkeypatch.setattr(collage, "build_image_grid", fake_grid)
    monkeypatch.setattr(collage, "_upload_bytes", fake_upload)

    result = await collage.render_image_collage(
        ["one", "two"], labels=["first", "second"], filename_prefix="test"
    )

    assert result == {"type": "url", "url": "https://cdn.example/collage.png"}
    assert captured["images"] == [image]
    assert captured["labels"] == ["first"]
    assert captured["filename"].startswith("test_")


@pytest.mark.asyncio
async def test_collage_falls_back_to_base64_when_upload_fails(monkeypatch):
    image = Image.new("RGB", (2, 2), "white")

    async def fake_load(urls):
        return [(urls[0], image)]

    async def failed_upload(image_bytes, filename):
        raise RuntimeError("offline")

    monkeypatch.setattr(collage, "load_images", fake_load)
    monkeypatch.setattr(collage, "build_image_grid", lambda **kwargs: image)
    monkeypatch.setattr(collage, "_upload_bytes", failed_upload)

    result = await collage.render_image_collage(["one"])

    assert result["type"] == "base64"
    assert result["media_type"] == "image/png"
    assert result["data"]


@pytest.mark.asyncio
@pytest.mark.parametrize(
    ("adapter", "args", "expected_prefix"),
    [
        (aigc_channel._build_images_collage, (["u"],), "collage_detail"),
        (x._build_images_collage, (["u"],), "x_detail_collage"),
    ],
)
async def test_detail_adapters_delegate_to_shared_collage(
    monkeypatch, adapter, args, expected_prefix
):
    module = __import__(adapter.__module__, fromlist=["render_image_collage"])
    captured = {}

    async def fake_render(urls, **kwargs):
        captured.update(kwargs)
        return {"type": "url", "url": "ok"}

    monkeypatch.setattr(module, "render_image_collage", fake_render)

    assert await adapter(*args) == {"type": "url", "url": "ok"}
    assert captured["filename_prefix"] == expected_prefix


@pytest.mark.asyncio
async def test_search_adapters_keep_platform_selection_rules(monkeypatch):
    calls = []

    async def fake_render(urls, **kwargs):
        calls.append((urls, kwargs))
        return {"type": "url", "url": "ok"}

    monkeypatch.setattr(x, "render_image_collage", fake_render)
    monkeypatch.setattr(youtube, "render_image_collage", fake_render)

    await x._build_tweet_collage(
        [{"image_url_list": [{"image_url": "x.jpg"}], "channel_account_name": "a"}]
    )
    await youtube._build_video_collage(
        [{"thumbnails": [{"url": "yt.jpg"}], "title": "video"}]
    )

    assert calls[0][0] == ["x.jpg"]
    assert calls[0][1]["labels"] == ["@a"]
    assert calls[0][1]["filename_prefix"] == "x_collage"
    assert calls[1][0] == ["yt.jpg"]
    assert calls[1][1]["labels"] == ["video"]
    assert calls[1][1]["filename_prefix"] == "youtube_collage"