Create-knowledge-find-decode/tests/test_media.py

"""媒体落盘离线测试：注入假 downloader，断言图片落盘 + card.url 改写为 /data 路径。"""
from __future__ import annotations

from creation_knowledge import media
from core.models import Card, Post

PNG = b"\x89PNG\r\n\x1a\n" + b"fakepng"
JPG = b"\xff\xd8\xff\xe0" + b"fakejpg"


def _post():
    return Post(
        id="xhs_1", platform="xiaohongshu", url="u", content_id="1",
        image_urls=["http://cdn/a", "http://cdn/b"],
        cards=[Card(index=1, kind="image", url="http://cdn/a"),
               Card(index=2, kind="image", url="http://cdn/b")],
    )


def test_media_target():
    t = media.media_target("data", "20260616-101010", "douyin", "dy_9")
    assert str(t.dir).endswith("data/20260616-101010/douyin/dy_9")
    assert t.public_base == "/data/20260616-101010/douyin/dy_9"


def test_save_images_writes_and_rewrites(tmp_path):
    bodies = {"http://cdn/a": PNG, "http://cdn/b": JPG}
    calls = []

    def dl(url, platform):
        calls.append((url, platform))
        return bodies[url]

    post = _post()
    media.save_images(post, tmp_path, "/data/B/xiaohongshu/xhs_1", downloader=dl)

    assert (tmp_path / "image_1.png").read_bytes() == PNG  # 文件头嗅探出 .png
    assert (tmp_path / "image_2.jpg").read_bytes() == JPG
    assert post.cards[0].url == "/data/B/xiaohongshu/xhs_1/image_1.png"
    assert post.cards[1].url == "/data/B/xiaohongshu/xhs_1/image_2.jpg"
    assert calls == [("http://cdn/a", "xiaohongshu"), ("http://cdn/b", "xiaohongshu")]


def test_save_images_skips_failures(tmp_path):
    """单张下载失败只跳过，不抛、不改这张卡 url。"""
    def dl(url, platform):
        if url.endswith("b"):
            raise RuntimeError("boom")
        return JPG

    post = _post()
    media.save_images(post, tmp_path, "/data/B/xiaohongshu/xhs_1", downloader=dl)

    assert post.cards[0].url == "/data/B/xiaohongshu/xhs_1/image_1.jpg"
    assert post.cards[1].url == "http://cdn/b"  # 失败的保持原样


def test_save_images_ignores_local_and_segment(tmp_path):
    """已是本地路径 / 非 image 卡不处理。"""
    post = Post(id="x", platform="douyin", url="u", content_id="1",
                cards=[Card(index=1, kind="segment", url="/data/x/video.mp4", start=0, end=3),
                       Card(index=2, kind="image", url="/data/x/image_2.jpg")])

    def dl(url, platform):
        raise AssertionError("不应被调用")

    media.save_images(post, tmp_path, "/data/x", downloader=dl)
    assert post.cards[0].url == "/data/x/video.mp4"
    assert post.cards[1].url == "/data/x/image_2.jpg"