"""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()