yunke/yunke-server/tests/test_scrape_moments.py

"""scrape_moments 归一化 + 翻页内核 + 图片本地化测试（不启浏览器）。

云客接口事实（2026-08-25 真机抓包，已验证）：
``POST /wechatweb/wechatFriendster/getFriendsterInfo``，form 小驼峰字段
``startTime``/``endTime``(空串)/``pageNum``(从1)/``pageSize``/``wechatId``；
信封 ``data.data[]`` + ``data.hashNext``；行的 ``createTime`` 是**秒**
（与消息接口的 ms 不同，归一化 ×1000）。
"""

from __future__ import annotations

import hashlib

import httpx
import respx
import sqlalchemy as sa
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine

from app.db.base import Base
from app.db.models import Media, MomentsPost
from app.scraper.media import download_moments_images
from app.session.yunke_client import YunkeClient
from app.storage.object_store import ObjectStore
from scripts.scrape_via_page import _iter_moments_pages, _normalize_moments

RAW = {  # 2026-08-25 真机抓包样本（截取关键字段）
    "id": "69e082a3c3524a7fe3ac53b4", "wechatId": "wxid_wo97u1u2x11w22",
    "snsId": "-3546042415332912522", "friendsterName": "Yoyo-HESU",
    "friendsterWXId": "wxid_wo97u1u2x11w22",
    "friendsterWXIdImg": "https://wx.qlogo.cn/mmhead/ver_1/vsHn.../0",
    "content": "Welcome to HESU", "type": "1", "createTime": "1776302051",  # 秒！
    "praise": [{"name": "E**U", "praiseId": "wxid***r822", "praiseImg": "https://wx.qlogo.cn/p/0"}],
    "comment": [],
    "imgUrlsList": ["https://wx.oss.yunkecn.com/friendstr/8m0atd/16/b1/ea/xx.jpg"],
}


# ---------------------------------------------------------------------------
# _normalize_moments
# ---------------------------------------------------------------------------


def test_normalize_moments_converts_seconds_to_ms():
    rec = _normalize_moments(RAW)
    assert rec["id"] == "69e082a3c3524a7fe3ac53b4"
    assert rec["create_time"] == 1776302051 * 1000  # 秒→ms
    assert rec["author_name"] == "Yoyo-HESU"
    assert rec["praise_json"][0]["name"] == "E**U"
    assert rec["image_urls"] == RAW["imgUrlsList"]
    assert rec["raw_json"]["snsId"] == "-3546042415332912522"


def test_normalize_moments_null_safe():
    rec = _normalize_moments({})
    assert rec["id"] is None  # 缺 id 行由调用方丢弃
    assert rec["create_time"] is None
    assert rec["praise_json"] == []
    assert rec["image_urls"] == []


def test_normalize_moments_bad_create_time_is_none():
    rec = _normalize_moments({**RAW, "createTime": "not-a-number"})
    assert rec["create_time"] is None
    assert rec["id"] == RAW["id"]  # 其余字段不受影响


# ---------------------------------------------------------------------------
# _iter_moments_pages（翻页内核，fetch 回调注入）
# ---------------------------------------------------------------------------


def _page(rows: list[dict], hash_next: bool) -> dict:
    return {"code": "10000", "data": {"hashNext": hash_next, "data": rows}}


async def _noop_sleep(_seconds: float) -> None:
    return None


async def test_iter_moments_pages_accumulates_two_pages_and_stops():
    """第一页 hashNext=True → 继续；第二页 hashNext=False → 停（恰好 2 次请求）。"""
    bodies: list[dict] = []

    async def fetch(body: dict) -> dict:
        bodies.append(body)
        return _page([RAW] if body["pageNum"] == 1 else [{**RAW, "id": "second"}], body["pageNum"] == 1)

    records = await _iter_moments_pages(
        fetch, "wxid_wo97u1u2x11w22", page_size=100, max_pages=50, sleep=_noop_sleep
    )

    assert [r["id"] for r in records] == [RAW["id"], "second"]
    assert records[0]["create_time"] == 1776302051 * 1000
    # 恰好两次请求：hashNext=False 即止
    assert len(bodies) == 2
    # form 字段与真机抓包一致（小驼峰，pageNum 从 1 起）
    assert bodies[0] == {
        "startTime": "", "endTime": "", "pageNum": 1, "pageSize": 100,
        "wechatId": "wxid_wo97u1u2x11w22",
    }
    assert bodies[1]["pageNum"] == 2


async def test_iter_moments_pages_stops_when_hash_next_falsy():
    """单页 hashNext=False → 只请求一次。"""
    calls: list[int] = []

    async def fetch(body: dict) -> dict:
        calls.append(body["pageNum"])
        return _page([RAW], hash_next=False)

    records = await _iter_moments_pages(fetch, "wx", sleep=_noop_sleep)

    assert len(records) == 1
    assert calls == [1]


async def test_iter_moments_pages_respects_max_pages_cap():
    """hashNext 永真 → max_pages 页上限防死循环（4h 调度浅扫 max-pages 3 靠它）。"""
    calls: list[int] = []

    async def fetch(body: dict) -> dict:
        calls.append(body["pageNum"])
        return _page([{**RAW, "id": f"id-{body['pageNum']}"}], hash_next=True)

    records = await _iter_moments_pages(fetch, "wx", max_pages=3, sleep=_noop_sleep)

    assert len(calls) == 3
    assert [r["id"] for r in records] == ["id-1", "id-2", "id-3"]


async def test_iter_moments_pages_drops_null_id_rows_and_bad_envelope():
    """缺 id 行丢弃；data 非 dict / data.data 非列表 → 停止不抛。"""
    async def fetch(body: dict) -> dict:
        if body["pageNum"] == 1:
            return _page([RAW, {"content": "no id"}], hash_next=True)
        return _page([], hash_next=False)

    records = await _iter_moments_pages(fetch, "wx", sleep=_noop_sleep)
    assert [r["id"] for r in records] == [RAW["id"]]

    async def fetch_bad(_body: dict) -> dict:
        return {"code": "10000", "data": None}

    assert await _iter_moments_pages(fetch_bad, "wx", sleep=_noop_sleep) == []


# ---------------------------------------------------------------------------
# download_moments_images（朋友圈图片本地化）
# ---------------------------------------------------------------------------

BASE = "https://crm2.yunkecn.com"
OSS_A = "https://wx.oss.yunkecn.com/friendstr/8m0atd/16/b1/ea/a.jpg"
OSS_B = "https://wx.oss.yunkecn.com/friendstr/8m0atd/16/b1/ea/b.jpg"


async def _make_engine():
    engine = create_async_engine("sqlite+aiosqlite:///:memory:")
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
    return engine


def _primed_client() -> YunkeClient:
    return YunkeClient(base_url=BASE, userid="a" * 32, company="8m0atd")


@respx.mock
async def test_download_moments_images_localizes_all_urls(tmp_path):
    """一帖两图(image_urls JSON 数组)全部下载入 media 表,media_type='image'。"""
    store = ObjectStore(tmp_path)
    engine = await _make_engine()
    SessionLocal = async_sessionmaker(engine, expire_on_commit=False)

    blob_a, blob_b = b"\xff\xd8\xff\xe0jpg-a", b"\xff\xd8\xff\xe0jpg-b"
    respx.get(OSS_A).mock(return_value=httpx.Response(200, content=blob_a))
    respx.get(OSS_B).mock(return_value=httpx.Response(200, content=blob_b))

    client = _primed_client()
    try:
        async with SessionLocal() as session:
            session.add(MomentsPost(
                id="p1", wechat_id="wx1", image_urls=[OSS_A, OSS_B],
            ))
            await session.commit()

        async with SessionLocal() as session:
            attempted, localized = await download_moments_images(client, store, session)
            await session.commit()

        assert (attempted, localized) == (2, 2)
        for blob, url in ((blob_a, OSS_A), (blob_b, OSS_B)):
            sha = hashlib.sha256(blob).hexdigest()
            assert (tmp_path / sha[:2] / f"{sha}.jpg").is_file()
            async with SessionLocal() as session:
                media = (await session.execute(
                    sa.select(Media).where(Media.media_hash == sha)
                )).scalar_one()
                assert media.media_type == "image"
                assert media.channel == "wechat"
                assert media.original_url == url
                assert media.mime_type == "image/jpeg"
    finally:
        await client.aclose()
        await engine.dispose()


@respx.mock
async def test_download_moments_images_idempotent_via_original_url(tmp_path):
    """MomentsPost 无 media_hash 回链列 → 已有同 original_url 的 Media 行即跳过
    (第二轮 attempted=0,不触网)。"""
    store = ObjectStore(tmp_path)
    engine = await _make_engine()
    SessionLocal = async_sessionmaker(engine, expire_on_commit=False)

    respx.get(OSS_A).mock(return_value=httpx.Response(200, content=b"blob"))

    client = _primed_client()
    try:
        async with SessionLocal() as session:
            session.add(MomentsPost(
                id="p1", wechat_id="wx1", image_urls=[OSS_A],
            ))
            await session.commit()

        async with SessionLocal() as session:
            first = await download_moments_images(client, store, session)
            await session.commit()
        assert first == (1, 1)

        async with SessionLocal() as session:
            second = await download_moments_images(client, store, session)
            await session.commit()
        assert second == (0, 0)
    finally:
        await client.aclose()
        await engine.dispose()


@respx.mock
async def test_download_moments_images_skips_non_http_and_survives_failure(tmp_path):
    """非 http 条目跳过;下载 500 只计数不抛(下轮重试)。"""
    store = ObjectStore(tmp_path)
    engine = await _make_engine()
    SessionLocal = async_sessionmaker(engine, expire_on_commit=False)

    respx.get(OSS_A).mock(return_value=httpx.Response(500))

    client = _primed_client()
    try:
        async with SessionLocal() as session:
            session.add(MomentsPost(
                id="p1", wechat_id="wx1", image_urls=[OSS_A, "not-a-url"],
            ))
            await session.commit()

        async with SessionLocal() as session:
            attempted, localized = await download_moments_images(client, store, session)
            await session.commit()

        assert (attempted, localized) == (1, 0)
        async with SessionLocal() as session:
            assert (await session.execute(sa.select(Media))).scalars().all() == []
    finally:
        await client.aclose()
        await engine.dispose()