yunke/yunke-server/tests/test_message_images_downloader.py

"""Tests for :func:`app.scraper.media.download_message_images`.

Covers the message-media batch downloader that localizes chat-message
``file_url`` / ``file_th_url`` / ``file_hd_url`` OSS URLs into the local
object store, writes a ``Media`` row per blob, and links the message back via
``media_hash``.

Mirrors ``test_recording_downloader.py``:

* Happy path: one image message -> thumbnail (``file_th_url``) is downloaded,
  stored, upserted, linked.
* Thumbnail preference: image with both ``file_url`` and ``file_th_url`` set
  -> the thumbnail URL is fetched (not the original).
* Idempotent: a message already carrying ``media_hash`` is excluded from the
  candidate set (``attempted == 0``).
* Resilient: a persistent 500 does NOT hang the loop (the per-run ``tried``
  set excludes the failing msg_id from later batches) and returns
  ``(attempted=1, localized=0)`` with no exception escaping.
* Filter: a text message with an http URL (e.g. a documentation link) is
  NOT selected (only ``type IN ('image','image_legacy','video','file')`` rows
  are eligible).
* ``image_legacy`` coverage: legacy WeChat type-3 rows whose ``file_url``
  carries a valid OSS URL ARE selected and stored as ``media_type='image'``.
"""

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, Message
from app.scraper.media import download_message_images
from app.session.yunke_client import YunkeClient
from app.storage.object_store import ObjectStore

BASE = "https://crm2.yunkecn.com"
# Realistic OSS image URL shapes (yunke OSS uses ``wx.oss.yunkecn.com``; the
# downloader is URL-agnostic — any http(s) URL works through the mocked
# transport). Thumbnail + original + hd variants of the same asset.
OSS_THUMB = "https://wx.oss.yunkecn.com/wechat/img/abc_thumb.jpg"
OSS_FULL = "https://wx.oss.yunkecn.com/wechat/img/abc.jpg"
OSS_HD = "https://wx.oss.yunkecn.com/wechat/img/abc_hd.jpg"

USERID = "a" * 32
COMPANY = "8m0atd"


# ---------------------------------------------------------------------------
# Shared helpers (mirror tests/test_recording_downloader.py)
# ---------------------------------------------------------------------------


def _primed_client() -> YunkeClient:
    """A YunkeClient already in a logged-in state (login flow skipped)."""
    return YunkeClient(base_url=BASE, userid=USERID, company=COMPANY)


async def _make_engine():
    """Create an in-memory async SQLite engine with all tables created."""
    engine = create_async_engine("sqlite+aiosqlite:///:memory:")
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
    return engine


# ---------------------------------------------------------------------------
# download_message_images
# ---------------------------------------------------------------------------


@respx.mock
async def test_download_message_images_localizes_and_links(tmp_path):
    """One image message with an OSS thumb url: downloaded, stored, linked.

    Returns ``(1, 1)``; the message's ``media_hash`` is the sha256 of the
    fetched bytes; a ``Media`` row exists with ``media_type='image'``,
    ``mime_type='image/jpeg'``, ``original_url == OSS_FULL``; the file exists
    under ``tmp_path`` at the sha256-content-addressed ``.jpg`` path.
    """
    store = ObjectStore(tmp_path)
    engine = await _make_engine()
    SessionLocal = async_sessionmaker(engine, expire_on_commit=False)

    blob = b"\xff\xd8\xff\xe0thumb-jpg-payload"
    respx.get(OSS_THUMB).mock(return_value=httpx.Response(200, content=blob))

    client = _primed_client()
    try:
        async with SessionLocal() as session:
            session.add(
                Message(
                    msg_id="m1",
                    channel="wechat",
                    type="image",
                    file_url=OSS_FULL,
                    file_th_url=OSS_THUMB,
                    file_hd_url=OSS_HD,
                )
            )
            await session.commit()

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

        assert (attempted, localized) == (1, 1)

        # File landed on disk at the sha256-content-addressed path with .jpg.
        expected_sha = hashlib.sha256(blob).hexdigest()
        expected_path = tmp_path / expected_sha[:2] / f"{expected_sha}.jpg"
        assert expected_path.is_file()
        assert expected_path.read_bytes() == blob

        async with SessionLocal() as session:
            msg = (
                await session.execute(
                    sa.select(Message).where(Message.msg_id == "m1")
                )
            ).scalar_one()
            assert msg.media_hash == expected_sha

            media = (
                await session.execute(
                    sa.select(Media).where(Media.media_hash == expected_sha)
                )
            ).scalar_one()
            assert media.media_type == "image"
            assert media.mime_type == "image/jpeg"
            # All three OSS URLs are preserved on the Media row.
            assert media.original_url == OSS_FULL
            assert media.original_th_url == OSS_THUMB
            assert media.original_hd_url == OSS_HD
            assert media.local_path == str(expected_path)
            assert media.size_bytes == len(blob)
    finally:
        await client.aclose()
        await engine.dispose()


@respx.mock
async def test_image_prefers_thumbnail_url(tmp_path):
    """For ``type='image'`` rows, the THUMBNAIL URL is fetched first.

    The row carries all three URLs (file_url / file_th_url / file_hd_url); the
    downloader fetches ``file_th_url`` (smaller, sufficient for chat-timeline
    display per spec). Verified by asserting only the thumb route was hit.
    """
    store = ObjectStore(tmp_path)
    engine = await _make_engine()
    SessionLocal = async_sessionmaker(engine, expire_on_commit=False)

    thumb_blob = b"thumb-bytes"
    full_route = respx.get(OSS_FULL).mock(return_value=httpx.Response(200, content=b"full"))
    thumb_route = respx.get(OSS_THUMB).mock(return_value=httpx.Response(200, content=thumb_blob))

    client = _primed_client()
    try:
        async with SessionLocal() as session:
            session.add(
                Message(
                    msg_id="m-pref",
                    channel="wechat",
                    type="image",
                    file_url=OSS_FULL,
                    file_th_url=OSS_THUMB,
                )
            )
            await session.commit()

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

        assert (attempted, localized) == (1, 1)
        assert thumb_route.called
        assert not full_route.called

        # media row is keyed by the THUMBNAIL's sha256.
        expected_sha = hashlib.sha256(thumb_blob).hexdigest()
        async with SessionLocal() as session:
            msg = (
                await session.execute(
                    sa.select(Message).where(Message.msg_id == "m-pref")
                )
            ).scalar_one()
            assert msg.media_hash == expected_sha
    finally:
        await client.aclose()
        await engine.dispose()


@respx.mock
async def test_video_uses_file_url(tmp_path):
    """For ``type='video'`` rows, ``file_url`` (the original) is fetched.

    The image-thumbnail preference does NOT apply to video; the original
    ``file_url`` is the right asset.
    """
    store = ObjectStore(tmp_path)
    engine = await _make_engine()
    SessionLocal = async_sessionmaker(engine, expire_on_commit=False)

    blob = b"\x00\x00\x00 mp4-isom-payload"
    video_url = "https://wx.oss.yunkecn.com/wechat/vid/abc.mp4"
    respx.get(video_url).mock(return_value=httpx.Response(200, content=blob))

    client = _primed_client()
    try:
        async with SessionLocal() as session:
            session.add(
                Message(
                    msg_id="m-vid",
                    channel="wechat",
                    type="video",
                    file_url=video_url,
                )
            )
            await session.commit()

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

        assert (attempted, localized) == (1, 1)
        expected_sha = hashlib.sha256(blob).hexdigest()
        expected_path = tmp_path / expected_sha[:2] / f"{expected_sha}.mp4"
        assert expected_path.is_file()

        async with SessionLocal() as session:
            msg = (
                await session.execute(
                    sa.select(Message).where(Message.msg_id == "m-vid")
                )
            ).scalar_one()
            assert msg.media_hash == expected_sha
    finally:
        await client.aclose()
        await engine.dispose()


@respx.mock
async def test_idempotent_skip(tmp_path):
    """A message already carrying ``media_hash`` is excluded -> attempted 0.

    The WHERE clause filters ``media_hash IS NULL``, so a previously localized
    message is never re-fetched.
    """
    store = ObjectStore(tmp_path)
    engine = await _make_engine()
    SessionLocal = async_sessionmaker(engine, expire_on_commit=False)

    client = _primed_client()
    try:
        async with SessionLocal() as session:
            session.add(
                Message(
                    msg_id="m2",
                    channel="wechat",
                    type="image",
                    file_url=OSS_FULL,
                    file_th_url=OSS_THUMB,
                    media_hash="already",
                )
            )
            await session.commit()

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

        assert (attempted, localized) == (0, 0)
        # No network was touched: the candidate set was empty.
        assert respx.calls.call_count == 0
    finally:
        await client.aclose()
        await engine.dispose()


@respx.mock
async def test_resilient_on_failure_and_no_infinite_loop(tmp_path):
    """A persistent 500: no exception, loop TERMINATES, ``(1, 0)`` returned.

    This is the key guard for the per-run ``tried`` set: without it, the
    failing row would stay ``media_hash IS NULL`` and be re-selected by every
    subsequent batch forever, hanging the run. The ``tried`` set excludes
    the attempted msg_id from later batches, so the next SELECT returns an
    empty page and the loop breaks.
    """
    store = ObjectStore(tmp_path)
    engine = await _make_engine()
    SessionLocal = async_sessionmaker(engine, expire_on_commit=False)

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

    client = _primed_client()
    try:
        async with SessionLocal() as session:
            session.add(
                Message(
                    msg_id="m3",
                    channel="wechat",
                    type="image",
                    file_th_url=OSS_THUMB,
                )
            )
            await session.commit()

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

        assert (attempted, localized) == (1, 0)

        # The failing message is left un-localized (retried on the NEXT run).
        async with SessionLocal() as session:
            msg = (
                await session.execute(
                    sa.select(Message).where(Message.msg_id == "m3")
                )
            ).scalar_one()
            assert msg.media_hash is None
    finally:
        await client.aclose()
        await engine.dispose()


@respx.mock
async def test_text_message_with_url_is_not_selected(tmp_path):
    """A text message whose text happens to be a URL is NOT selected.

    The candidate filter requires ``type IN ('image','image_legacy','video',
    'file')``; a text message is skipped even if its columns somehow carry an
    http URL.
    """
    store = ObjectStore(tmp_path)
    engine = await _make_engine()
    SessionLocal = async_sessionmaker(engine, expire_on_commit=False)

    client = _primed_client()
    try:
        async with SessionLocal() as session:
            session.add(
                Message(
                    msg_id="m-text",
                    channel="wechat",
                    type="text",
                    text="see https://wx.oss.yunkecn.com/wechat/img/abc.jpg",
                    # Even with a file_url, type='text' disqualifies the row.
                    file_url=OSS_FULL,
                )
            )
            await session.commit()

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

        assert (attempted, localized) == (0, 0)
        assert respx.calls.call_count == 0
    finally:
        await client.aclose()
        await engine.dispose()


@respx.mock
async def test_image_legacy_is_selected_and_normalized_to_image(tmp_path):
    """``type='image_legacy'`` rows ARE selected (smoke.db 26 rows).

    The candidate filter includes ``image_legacy`` so the legacy WeChat type-3
    rows whose ``file_url`` carries a valid ``wx.oss.yunkecn.com`` OSS URL get
    localized. The thumbnail-first preference applies (same as ``image``), and
    the stored ``Media.media_type`` is NORMALIZED to ``'image'`` so the legacy
    bucket lands in the same bin as type-2 images.
    """
    store = ObjectStore(tmp_path)
    engine = await _make_engine()
    SessionLocal = async_sessionmaker(engine, expire_on_commit=False)

    thumb_blob = b"\xff\xd8\xff\xe0legacy-thumb-jpg"
    thumb_route = respx.get(OSS_THUMB).mock(
        return_value=httpx.Response(200, content=thumb_blob)
    )
    full_route = respx.get(OSS_FULL).mock(return_value=httpx.Response(200, content=b"full"))

    client = _primed_client()
    try:
        async with SessionLocal() as session:
            session.add(
                Message(
                    msg_id="m-legacy",
                    channel="wechat",
                    type="image_legacy",
                    type_raw=3,
                    file_url=OSS_FULL,
                    file_th_url=OSS_THUMB,
                    file_hd_url=OSS_HD,
                )
            )
            await session.commit()

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

        assert (attempted, localized) == (1, 1)
        # Thumbnail-first preference applies to image_legacy exactly as for image.
        assert thumb_route.called
        assert not full_route.called

        expected_sha = hashlib.sha256(thumb_blob).hexdigest()
        async with SessionLocal() as session:
            msg = (
                await session.execute(
                    sa.select(Message).where(Message.msg_id == "m-legacy")
                )
            ).scalar_one()
            assert msg.media_hash == expected_sha

            media = (
                await session.execute(
                    sa.select(Media).where(Media.media_hash == expected_sha)
                )
            ).scalar_one()
            # Normalized: image_legacy -> image on the stored Media row.
            assert media.media_type == "image"
            assert media.mime_type == "image/jpeg"
            assert media.original_url == OSS_FULL
    finally:
        await client.aclose()
        await engine.dispose()