Blame for tests/test_ethumb.py

"""/ethumb: ask for a file or video, send back its thumbnail."""
import io
import shutil
import time
from pathlib import Path
from types import SimpleNamespace

from PIL import Image
from telethon import types

from bot import app as appmod

from .appkit import OWNER, STRANGER, auto_thumb, click, data_of, jpeg, say, send_doc
from .conftest import needs_ffmpeg
from .tgfakes import FakeEvent, FakeMessage, dispatch, document_message

REAL_THUMB = [SimpleNamespace(type="m")]  # stands in for a PhotoSize
STRIPPED = [types.PhotoStrippedSize(type="i", bytes=b"\x01\x14\x28abc")]


def photos_sent(ctx):
    return [(entity, file, kw) for entity, file, kw in ctx.client.files if isinstance(file, io.BytesIO)]


def temp_leftovers(ctx):
    root = Path(ctx.cfg.temp_dir)
    return list(root.glob("ethumb*")) if root.exists() else []


def serve(monkeypatch, media_dir, name="sample.mkv"):
    """Make transfer.download copy a real sample file; returns the list of calls."""
    calls = []

    async def fake_download(client, message, dest, **kw):
        calls.append(kw)
        shutil.copy(media_dir / name, dest)

    monkeypatch.setattr(appmod.transfer, "download", fake_download)
    return calls


# ---------------------------------------------------------------------------
# the flow
# ---------------------------------------------------------------------------


async def test_ethumb_asks_for_a_file(ctx):
    ev = await say(ctx, OWNER, "/ethumb")
    assert "EXTRACT THUMBNAIL" in ev.texts[0] and "one file at a time" in ev.texts[0] and "/cancel" in ev.texts[0]
    assert data_of(ev.out[0][1]) == ["et:stop"]
    mode = ctx.app.modes[OWNER]
    assert mode.kind == "ethumb" and mode.expires > time.time() + 500


async def test_the_files_own_telegram_thumbnail_is_sent_when_it_has_one(ctx, monkeypatch):
    async def no_download(*a, **kw):
        raise AssertionError("nothing needs downloading when Telegram already has a preview")

    monkeypatch.setattr(appmod.transfer, "download", no_download)
    await say(ctx, OWNER, "/ethumb")
    ctx.client.thumb_data[60] = jpeg((320, 180))
    ev = await send_doc(ctx, OWNER, "Episode 01.mp4", thumbs=REAL_THUMB, mid=60)
    assert ev.texts == ["🖼 Extracting the thumbnail…"]
    (entity, file, kw), = photos_sent(ctx)
    assert entity == OWNER and file.name == "thumbnail.jpg" and Image.open(io.BytesIO(file.getvalue())).size == (320, 180)
    assert "Thumbnail of <b>Episode 01.mp4</b>" in kw["caption"] and kw["parse_mode"] == "html"
    assert (OWNER, [1001]) in ctx.client.deleted  # the "Extracting…" message goes
    assert OWNER not in ctx.app.modes and ctx.client.downloads == [(60, bytes, -1)]


@needs_ffmpeg
async def test_without_a_telegram_thumbnail_a_frame_is_grabbed_from_the_video(ctx, media_dir, monkeypatch):
    calls = serve(monkeypatch, media_dir)
    await say(ctx, OWNER, "/ethumb")
    await send_doc(ctx, OWNER, "Show - 05.mkv", size=(media_dir / "sample.mkv").stat().st_size, mid=61)
    (_e, file, kw), = photos_sent(ctx)
    img = Image.open(io.BytesIO(file.getvalue()))
    assert img.format == "JPEG" and img.size == (320, 240)
    assert len(calls) == 1 and calls[0]["max_bytes"] == appmod.TRACKS_PROBE_BYTES
    assert "Show - 05.mkv" in kw["caption"] and (OWNER, [1001]) in ctx.client.deleted
    assert temp_leftovers(ctx) == []


@needs_ffmpeg
async def test_the_whole_file_is_fetched_when_the_start_has_no_picture(ctx, media_dir, monkeypatch):
    calls = []

    async def fake_download(client, message, dest, **kw):
        calls.append(kw)
        if "max_bytes" in kw:
            Path(dest).write_bytes(b"\x00" * 4096)  # an mp4 whose index is at the end
        else:
            shutil.copy(media_dir / "sample.mp4", dest)

    monkeypatch.setattr(appmod.transfer, "download", fake_download)
    await say(ctx, OWNER, "/ethumb")
    await send_doc(ctx, OWNER, "Movie.mp4", size=40 * 1024 * 1024, mid=62)
    assert len(calls) == 2 and "max_bytes" not in calls[1]
    assert len(photos_sent(ctx)) == 1 and temp_leftovers(ctx) == []


@needs_ffmpeg
async def test_a_file_with_no_picture_anywhere_gets_an_honest_answer(ctx, monkeypatch):
    async def junk(client, message, dest, **kw):
        Path(dest).write_bytes(b"\x00" * 4096)

    monkeypatch.setattr(appmod.transfer, "download", junk)
    for size in (2 * 1024 * 1024, 900 * 1024 * 1024):  # small / too big to fetch whole
        ctx.client.edited.clear()
        await say(ctx, OWNER, "/ethumb")
        await send_doc(ctx, OWNER, "Broken.mkv", size=size, mid=63)
        assert "couldn't find a thumbnail" in ctx.client.edited[-1][1], size
    assert photos_sent(ctx) == [] and temp_leftovers(ctx) == []


async def test_files_that_cannot_have_a_frame_are_not_downloaded(ctx, monkeypatch):
    async def no_download(*a, **kw):
        raise AssertionError("a zip has no picture to grab")

    monkeypatch.setattr(appmod.transfer, "download", no_download)
    await say(ctx, OWNER, "/ethumb")
    await send_doc(ctx, OWNER, "archive.zip", size=5 * 1024 * 1024, video=False, mime="application/zip", mid=64)
    assert "couldn't find a thumbnail" in ctx.client.edited[-1][1] and photos_sent(ctx) == []


@needs_ffmpeg
async def test_the_tiny_blurred_preview_is_only_the_last_resort(ctx, media_dir, monkeypatch):
    serve(monkeypatch, media_dir)
    await say(ctx, OWNER, "/ethumb")
    ctx.client.thumb_data[65] = b"BLURRY-PREVIEW"
    await send_doc(ctx, OWNER, "Show.mkv", size=(media_dir / "sample.mkv").stat().st_size, thumbs=STRIPPED, mid=65)
    (_e, file, _kw), = photos_sent(ctx)
    assert file.getvalue() != b"BLURRY-PREVIEW" and Image.open(io.BytesIO(file.getvalue())).size == (320, 240)  # a real frame wins

    async def junk(client, message, dest, **kw):
        Path(dest).write_bytes(b"\x00" * 4096)

    monkeypatch.setattr(appmod.transfer, "download", junk)
    await say(ctx, OWNER, "/ethumb")
    await send_doc(ctx, OWNER, "Show2.mkv", size=2 * 1024 * 1024, thumbs=STRIPPED, mid=65)
    assert photos_sent(ctx)[-1][1].getvalue() == b"BLURRY-PREVIEW"  # nothing better could be made


async def test_ethumb_as_a_reply_works_at_once_without_asking(ctx):
    ctx.client.thumb_data[66] = jpeg((320, 180))
    reply = document_message("Episode.mp4", mid=66, thumbs=REAL_THUMB)
    await say(ctx, OWNER, "/ethumb", reply=reply, mid=67)
    assert len(photos_sent(ctx)) == 1 and OWNER not in ctx.app.modes


async def test_it_handles_one_file_per_command(ctx):
    ctx.client.thumb_data[70] = jpeg()
    await say(ctx, OWNER, "/ethumb")
    await send_doc(ctx, OWNER, "One.mp4", thumbs=REAL_THUMB, mid=70)
    extra = await send_doc(ctx, OWNER, "Two.mp4", thumbs=REAL_THUMB, mid=71)  # right after: a hint, not an auto-rename
    assert "/ethumb handles one file at a time" in extra.texts[0] and ctx.app.queue.waiting == 0
    assert (await send_doc(ctx, OWNER, "Three.mp4", thumbs=REAL_THUMB, mid=72)).out == []  # further ones: silence
    assert len(photos_sent(ctx)) == 1
    ctx.app._recent[OWNER][1] -= appmod.ONE_SHOT_GRACE + 1
    later = await send_doc(ctx, OWNER, "Four.mp4", mid=73)
    assert "Queued" in later.texts[0]  # normal auto-rename is back


async def test_a_photo_while_waiting_is_ignored_and_the_wait_goes_on(ctx):
    await say(ctx, OWNER, "/ethumb")
    ev = FakeEvent(OWNER, message=FakeMessage("", photo=True, mid=80), photo_bytes=jpeg())
    await dispatch(ctx.client, ev)
    assert ev.out == [] and await auto_thumb(ctx) is None and ctx.app.modes[OWNER].kind == "ethumb"  # a photo is nobody's thumbnail


async def test_cancel_and_the_stop_button_end_it(ctx):
    await say(ctx, OWNER, "/ethumb")
    ev = await say(ctx, OWNER, "/cancel")
    assert "Cancelled" in ev.texts[0] and OWNER not in ctx.app.modes
    ev = await send_doc(ctx, OWNER, "A.mkv", mid=90)
    assert "Queued" in ev.texts[0]

    await say(ctx, OWNER, "/ethumb")
    cb = await click(ctx, OWNER, "et:stop")
    assert "Stopped" in cb.edits[-1][0] and OWNER not in ctx.app.modes


async def test_it_expires(ctx):
    await say(ctx, OWNER, "/ethumb")
    ctx.app.modes[OWNER].expires = time.time() - 1
    ev = await send_doc(ctx, OWNER, "A.mkv", mid=91)
    assert "Queued" in ev.texts[0]


async def test_strangers_cannot_use_it(ctx):
    ev = await say(ctx, STRANGER, "/ethumb")
    assert ev.out == [] and STRANGER not in ctx.app.modes


async def test_starting_it_ends_a_sequence_and_says_so(ctx):
    await say(ctx, OWNER, "/sequence")
    await send_doc(ctx, OWNER, "Show - 01.mkv", mid=92)
    ev = await say(ctx, OWNER, "/ethumb")
    assert "1 file(s) was discarded" in ev.texts[0] and ctx.app.modes[OWNER].kind == "ethumb"


async def test_a_failing_extraction_is_reported_not_raised(ctx, monkeypatch):
    async def boom(*a, **kw):
        raise RuntimeError("network down")

    monkeypatch.setattr(appmod.transfer, "download", boom)
    await say(ctx, OWNER, "/ethumb")
    ev = await send_doc(ctx, OWNER, "A.mkv", mid=93)
    assert ev.texts == ["🖼 Extracting the thumbnail…"] and "couldn't find a thumbnail" in ctx.client.edited[-1][1]
    assert temp_leftovers(ctx) == []