HeXi/tests/test_douyin_web.py

"""抖音免浏览器链路（douyin_web.py）单元测试

覆盖：作品 id 提取 / 签名查询串 / 网关头构造 / 最高码率档选择 / 图集下载项收集。
签名与媒体选择都是纯函数，不碰网络（真实链路实测见 debug 目录脚本）。

douyin_web.py 有插件内相对导入（`...models` / `...utils` / `.douyin_sign` 等），
按 test_kuaishou_parse.py 同款姿势：sys.modules 预置包替身 + 真实加载依赖后裸加载。
"""

import importlib.util
import sys
import types
from pathlib import Path

import pytest

_REPO = Path(__file__).resolve().parents[1]
_PLUGIN = _REPO / "hexi" / "plugins" / "nonebot_plugin_video_analysis"
_PKG = "hexi.plugins.nonebot_plugin_video_analysis"

_REAL_MODULES = {
    f"{_PKG}.models": _PLUGIN / "models.py",
    f"{_PKG}.utils": _PLUGIN / "utils.py",
    f"{_PKG}.services.fetchers.douyin_parser": _PLUGIN
    / "services/fetchers/douyin_parser.py",
    f"{_PKG}.services.fetchers.douyin_sign": _PLUGIN
    / "services/fetchers/douyin_sign.py",
    f"{_PKG}.services.fetchers.douyin_api": _PLUGIN / "services/fetchers/douyin_api.py",
    f"{_PKG}.services.fetchers.douyin_web": _PLUGIN / "services/fetchers/douyin_web.py",
}
_STUB_PACKAGES = (_PKG, f"{_PKG}.services", f"{_PKG}.services.fetchers")


def _load_real(name: str, path: Path):
    spec = importlib.util.spec_from_file_location(name, path)
    module = importlib.util.module_from_spec(spec)
    sys.modules[name] = module
    spec.loader.exec_module(module)
    return module


@pytest.fixture(scope="module")
def dw():
    saved = {name: sys.modules.get(name) for name in _STUB_PACKAGES}
    for name in _STUB_PACKAGES:
        sys.modules[name] = types.ModuleType(name)
    for name, path in _REAL_MODULES.items():
        saved[name] = sys.modules.get(name)
        _load_real(name, path)

    yield sys.modules[f"{_PKG}.services.fetchers.douyin_web"]

    for name, previous in saved.items():
        if previous is None:
            sys.modules.pop(name, None)
        else:
            sys.modules[name] = previous


# ───────────────────────── 样本 ─────────────────────────


def _image(url_list=None, download_url_list=None, video=None) -> dict:
    img = {
        "uri": "img-uri",
        "width": 2160,
        "height": 2880,
        "url_list": (
            url_list
            if url_list is not None
            else ["https://p3.douyinpic.com/a~q75.webp"]
        ),
        "download_url_list": (
            download_url_list
            if download_url_list is not None
            else ["https://p9.douyinpic.com/a~tplv-dy-water-v2:xxx.webp"]
        ),
    }
    if video is not None:
        img["video"] = video
    return img


def _video(bit_rate=None, play_urls=("https://v95.douyinvod.com/full.mp4",)) -> dict:
    v = {"play_addr": {"url_list": list(play_urls), "width": 1080, "height": 1440}}
    if bit_rate is not None:
        v["bit_rate"] = bit_rate
    return v


def _gear(br, urls, gear_name="gear", h265=0, size=100):
    return {
        "gear_name": gear_name,
        "bit_rate": br,
        "is_h265": h265,
        "play_addr": {
            "url_list": list(urls),
            "width": 1080,
            "height": 1440,
            "data_size": size,
        },
    }


# ───────────────────────── extract_aweme_id ─────────────────────────


@pytest.mark.parametrize(
    "url,expected",
    [
        ("https://www.douyin.com/note/7696029938990839419", "7696029938990839419"),
        (
            "https://www.douyin.com/video/7618117951811949541?foo=1",
            "7618117951811949541",
        ),
        (
            "https://www.iesdouyin.com/share/video/7412345678901234567/?region=CN",
            "7412345678901234567",
        ),
        (
            "https://www.iesdouyin.com/share/note/7412345678901234567/",
            "7412345678901234567",
        ),
        ("https://v.douyin.com/abc123/", None),
        ("https://github.com/x/y", None),
    ],
)
def test_extract_aweme_id(dw, url, expected):
    assert dw.extract_aweme_id(url) == expected


# ───────────────────────── 签名 ─────────────────────────


def test_build_signed_query(dw):
    qs = dw.build_signed_query("7618117951811949541")
    assert "aweme_id=7618117951811949541" in qs
    assert "a_bogus=" in qs
    a_bogus = qs.split("a_bogus=")[-1]
    assert len(a_bogus) >= 100
    # 两次签名都有效（签名带随机熵，不要求逐字节一致）
    qs2 = dw.build_signed_query("7618117951811949541")
    assert "a_bogus=" in qs2


def test_build_signed_query_requires_aweme_id_last(dw):
    """aweme_id 应在查询串末尾、a_bogus 追加在其后（与浏览器请求形态一致）"""
    qs = dw.build_signed_query("123")
    assert qs.index("aweme_id=123") < qs.index("a_bogus=")


# ───────────────────────── 网关头 ─────────────────────────


def test_build_headers_always_argus_and_uifid(dw):
    headers = dw._build_headers({"UIFID": "u-123"})
    assert headers["x-tt-argus"] == "1"
    assert headers["uifid"] == "u-123"
    assert headers["User-Agent"] == dw.WEB_UA


def test_build_headers_uifid_temp_fallback(dw):
    headers = dw._build_headers({"UIFID_TEMP": "tmp"})
    assert headers["uifid"] == "tmp"
    assert "uifid" not in dw._build_headers({})


def test_cookie_map_filters_domain(dw):
    cookies = [
        {"name": "UIFID", "value": "u1", "domain": ".douyin.com"},
        {"name": "a1", "value": "x", "domain": ".xiaohongshu.com"},
        {"name": "ttwid", "value": "t", "domain": "www.douyin.com"},
    ]
    m = dw._douyin_cookie_map(cookies)
    assert m == {"UIFID": "u1", "ttwid": "t"}


# ───────────────────────── 最高码率档选择 ─────────────────────────


def test_best_gear_picks_max_bitrate(dw):
    video = {
        "bit_rate": [
            _gear(773939, ["https://cdn/1080.mp4"], gear_name="normal_1080_0"),
            _gear(385679, ["https://cdn/720.mp4"], gear_name="720_1_1", h265=1),
            _gear(
                1419088,
                ["https://cdn/best.mp4", "https://cdn/best-bak.mp4"],
                gear_name="1080_1_1",
            ),
        ],
        "play_addr": {"url_list": ["https://cdn/main.mp4"]},
    }
    urls, meta = dw._best_urls_from_video(video)
    assert urls == ["https://cdn/best.mp4", "https://cdn/best-bak.mp4"]
    assert meta["bit_rate"] == 1419088
    assert meta["width"] == 1080


def test_best_gear_falls_back_to_play_addr(dw):
    video = _video(bit_rate=[], play_urls=["https://cdn/only.mp4"])
    urls, meta = dw._best_urls_from_video(video)
    assert urls == ["https://cdn/only.mp4"]
    assert meta["gear_name"] is None


def test_best_gear_skips_gear_without_urls(dw):
    video = {
        "bit_rate": [
            {"bit_rate": 9999999, "play_addr": {"url_list": []}},
            _gear(100, ["https://cdn/ok.mp4"]),
        ],
        "play_addr": {"url_list": ["https://cdn/main.mp4"]},
    }
    urls, _ = dw._best_urls_from_video(video)
    assert urls == ["https://cdn/ok.mp4"]


# ───────────────────────── 下载项收集 ─────────────────────────


def test_collect_video_post(dw):
    detail = {
        "aweme_id": "1",
        "desc": "x",
        "video": _video(bit_rate=[_gear(500, ["https://cdn/v.mp4"])]),
    }
    items, is_img = dw.collect_download_items(detail)
    assert is_img is False
    assert items == [(["https://cdn/v.mp4"], ".mp4")]


def test_collect_video_post_without_urls_raises(dw):
    with pytest.raises(Exception):
        dw.collect_download_items({"video": {"play_addr": {"url_list": []}}})


def test_collect_images_prefers_no_watermark_urls(dw):
    detail = {
        "images": [
            _image(
                url_list=["https://p3/a-q75.webp", "https://p3/a-q75.jpeg"],
                download_url_list=["https://p9/a-water.webp"],
            )
        ]
    }
    items, is_img = dw.collect_download_items(detail)
    assert is_img is True
    assert items[0][0] == [
        "https://p3/a-q75.webp",
        "https://p3/a-q75.jpeg",  # 无水印在前
        "https://p9/a-water.webp",  # 带作者水印的保存版兜底
    ]
    assert items[0][1] == ""


def test_collect_animated_note_uses_video_gears(dw):
    anim_video = {
        "bit_rate": [_gear(300, ["https://cdn/anim-best.mp4"])],
        "play_addr": {"url_list": ["https://cdn/anim.mp4"]},
    }
    detail = {
        "images": [
            _image(url_list=[], download_url_list=[], video=anim_video),
            _image(url_list=[], download_url_list=[], video=anim_video),
        ]
    }
    items, is_img = dw.collect_download_items(detail)
    assert is_img is True
    assert items == [
        (["https://cdn/anim-best.mp4"], ".mp4"),
        (["https://cdn/anim-best.mp4"], ".mp4"),
    ]


def test_collect_mixed_note(dw):
    """图 + 内嵌视频混排：各自按类型收集"""
    detail = {
        "images": [
            _image(),
            _image(
                url_list=[],
                download_url_list=[],
                video={"play_addr": {"url_list": ["https://cdn/clip.mp4"]}},
            ),
        ]
    }
    items, is_img = dw.collect_download_items(detail)
    assert is_img is True
    assert len(items) == 2
    assert items[0][1] == ""  # 图片
    assert items[1] == (["https://cdn/clip.mp4"], ".mp4")


# ───────────────────────── 扩展名推断 ─────────────────────────


@pytest.mark.parametrize(
    "url,expected",
    [
        ("https://x/a.jpg?x=1", ".jpg"),
        ("https://x/a~q75.webp?y=2", ".webp"),
        ("https://x/a.jpeg", ".jpeg"),
        ("https://x/a", ".jpg"),
        ("https://x/a.heic", ".heic"),
    ],
)
def test_infer_extension(dw, url, expected):
    assert dw._infer_extension(url) == expected