"""/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 Episode 01.mp4" 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) == []