Spaces:
Running
Running
| """/resize: pick a size, send an image, get it back resized.""" | |
| import io | |
| import time | |
| from PIL import Image | |
| from bot import app as appmod | |
| from bot import panels | |
| from .appkit import OWNER, STRANGER, auto_thumb, autofmt, click, data_of, jpeg, say, send_doc | |
| from .tgfakes import FakeEvent, FakeMessage, dispatch | |
| FULL_HD = (1920, 1080) | |
| def results(ctx): | |
| return [(entity, file, kw) for entity, file, kw in ctx.client.files if isinstance(file, io.BytesIO)] | |
| def size_of(file): | |
| return Image.open(io.BytesIO(file.getvalue())).size | |
| async def send_photo(ctx, mid=50, data=None, sender=OWNER): | |
| ctx.client.media[mid] = data if data is not None else jpeg(FULL_HD) | |
| ev = FakeEvent(sender, message=FakeMessage("", photo=True, mid=mid), photo_bytes=jpeg()) | |
| await dispatch(ctx.client, ev) | |
| return ev | |
| async def send_image_file(ctx, name="poster.png", mid=51, data=None, size=200 * 1024): | |
| ctx.client.media[mid] = data if data is not None else jpeg(FULL_HD) | |
| return await send_doc(ctx, OWNER, name, size=size, video=False, mime="image/png", mid=mid) | |
| # --------------------------------------------------------------------------- | |
| # choosing the size | |
| # --------------------------------------------------------------------------- | |
| async def test_resize_offers_the_sizes(ctx): | |
| ev = await say(ctx, OWNER, "/resize") | |
| text, buttons = ev.texts[0], data_of(ev.out[0][1]) | |
| assert "RESIZE IMAGE" in text and "200 KB" in text and "/resize 320x180" in text | |
| assert buttons == ["rz:s:320", "rz:s:640", "rz:s:1280", "rz:custom", "x:close"] | |
| assert OWNER not in ctx.app.modes # nothing is waiting for an image until a size is chosen | |
| async def test_a_preset_button_makes_the_bot_wait_for_the_image(ctx): | |
| cb = await click(ctx, OWNER, "rz:s:320") | |
| assert "Size: <b>320 px</b>" in cb.edits[-1][0] and "one image at a time" in cb.edits[-1][0] | |
| assert data_of(cb.edits[-1][1]) == ["rz:menu", "rz:stop"] | |
| mode = ctx.app.modes[OWNER] | |
| assert mode.kind == "resize" and mode.data == {"box": (320, 320)} and mode.expires > time.time() + 500 | |
| await click(ctx, OWNER, "rz:s:1280") | |
| assert ctx.app.modes[OWNER].data["box"] == (1280, 1280) | |
| async def test_unknown_presets_are_ignored(ctx): | |
| cb = await click(ctx, OWNER, "rz:s:999") | |
| assert cb.edits == [] and OWNER not in ctx.app.modes | |
| async def test_the_size_can_be_given_with_the_command(ctx): | |
| ev = await say(ctx, OWNER, "/resize 320x180") | |
| assert "<b>320×180</b>" in ev.texts[0] and ctx.app.modes[OWNER].data["box"] == (320, 180) | |
| ev = await say(ctx, OWNER, "/resize 640") | |
| assert "<b>640 px</b>" in ev.texts[0] and ctx.app.modes[OWNER].data["box"] == (640, 640) | |
| for bad in ("tiny", "10x10", "99999", "320 px"): | |
| ev = await say(ctx, OWNER, f"/resize {bad}") | |
| assert "❌" in ev.texts[0] and "/resize 320x180" in ev.texts[0], bad | |
| assert ctx.app.modes[OWNER].data["box"] == (640, 640) # a bad size does not disturb the one that was set | |
| async def test_a_custom_size_is_typed_in(ctx): | |
| cb = await click(ctx, OWNER, "rz:custom", msg_id=500) | |
| p = ctx.app.pending[OWNER] | |
| assert p.kind == "resize:size" and p.msg_id == 500 and cb.edits[-1][0] == panels.RESIZE_CUSTOM_PROMPT | |
| assert data_of(cb.edits[-1][1]) == ["rz:menu"] | |
| ev = await say(ctx, OWNER, "nonsense", mid=60) | |
| assert "couldn't read that size" in ev.texts[0] and "Try again" in ev.texts[0] and OWNER in ctx.app.pending | |
| ev = await say(ctx, OWNER, "300 x 200", mid=61) | |
| assert ev.deleted and OWNER not in ctx.app.pending | |
| assert ctx.app.modes[OWNER].data["box"] == (300, 200) | |
| (chat, text, kw), = ctx.client.edited | |
| assert kw["message"] == 500 and "<b>300×200</b>" in text and data_of(kw["buttons"]) == ["rz:menu", "rz:stop"] | |
| async def test_cancelling_the_custom_prompt_goes_back_to_the_sizes(ctx): | |
| await click(ctx, OWNER, "rz:custom") | |
| cb = await click(ctx, OWNER, "rz:menu") | |
| assert OWNER not in ctx.app.pending and "RESIZE IMAGE" in cb.edits[-1][0] | |
| async def test_change_size_and_stop_buttons(ctx): | |
| await click(ctx, OWNER, "rz:s:320") | |
| cb = await click(ctx, OWNER, "rz:menu") | |
| assert "RESIZE IMAGE" in cb.edits[-1][0] and OWNER not in ctx.app.modes | |
| await click(ctx, OWNER, "rz:s:320") | |
| cb = await click(ctx, OWNER, "rz:stop") | |
| assert "Stopped" in cb.edits[-1][0] and OWNER not in ctx.app.modes | |
| # --------------------------------------------------------------------------- | |
| # resizing | |
| # --------------------------------------------------------------------------- | |
| async def test_a_photo_comes_back_resized(ctx): | |
| await say(ctx, OWNER, "/resize 320") | |
| await send_photo(ctx) | |
| (entity, file, kw), = results(ctx) | |
| assert entity == OWNER and file.name == "resized.jpg" and size_of(file) == (320, 180) | |
| assert len(file.getvalue()) <= 200 * 1024 | |
| assert "<b>320×180</b>" in kw["caption"] and "was 1920×1080" in kw["caption"] and kw["parse_mode"] == "html" | |
| assert kw["force_document"] is False | |
| assert OWNER not in ctx.app.modes | |
| assert ctx.client.downloads == [(50, bytes, None)] | |
| async def test_an_image_sent_as_a_file_works_too(ctx): | |
| await say(ctx, OWNER, "/resize 640x360") | |
| ev = await send_image_file(ctx) | |
| assert ev.out == [] and ctx.app.queue.waiting == 0 # not renamed | |
| (_e, file, _kw), = results(ctx) | |
| assert size_of(file) == (640, 360) | |
| async def test_what_is_not_an_image_does_not_use_up_the_request(ctx): | |
| await say(ctx, OWNER, "/resize 320") | |
| ev = await send_doc(ctx, OWNER, "Episode.mkv", mid=52) | |
| assert "isn't an image" in ev.texts[0] and ctx.app.queue.waiting == 0 and results(ctx) == [] | |
| assert ctx.app.modes[OWNER].kind == "resize" | |
| await send_photo(ctx) | |
| assert len(results(ctx)) == 1 | |
| async def test_it_handles_one_image_per_command_and_extra_photos_do_not_replace_the_thumbnail(ctx): | |
| await say(ctx, OWNER, "/resize 320") | |
| await send_photo(ctx, mid=50) | |
| extra = await send_photo(ctx, mid=51) | |
| assert "/resize handles one file at a time" in extra.texts[0] | |
| assert (await send_photo(ctx, mid=52)).out == [] | |
| assert len(results(ctx)) == 1 | |
| assert await autofmt(ctx) is None # none of the extra photos became a thumbnail (that would have made the auto format) | |
| doc = await send_image_file(ctx, mid=53) | |
| assert doc.out == [] and ctx.app.queue.waiting == 0 # an image file right after is skipped as well, not auto-renamed | |
| ctx.app._recent[OWNER][1] -= appmod.ONE_SHOT_GRACE + 1 | |
| await send_photo(ctx, mid=54) | |
| assert await auto_thumb(ctx) # after the grace a photo is the thumbnail of the auto format again | |
| async def test_large_results_are_sent_as_files_so_telegram_does_not_shrink_them(ctx): | |
| await say(ctx, OWNER, "/resize 2000") | |
| await send_photo(ctx, data=jpeg((3000, 2000))) | |
| (_e, file, kw), = results(ctx) | |
| assert size_of(file) == (2000, 1333) and kw["force_document"] is True | |
| async def test_resize_as_a_reply_works_at_once(ctx): | |
| ctx.client.media[40] = jpeg(FULL_HD) | |
| reply = FakeMessage("", photo=True, mid=40) | |
| await say(ctx, OWNER, "/resize", reply=reply, mid=41) | |
| (_e, file, _kw), = results(ctx) | |
| assert size_of(file) == (320, 180) and OWNER not in ctx.app.modes # thumbnail size unless told otherwise | |
| await say(ctx, OWNER, "/resize 640x640", reply=reply, mid=42) | |
| assert size_of(results(ctx)[-1][1]) == (640, 360) | |
| async def test_a_picture_that_cannot_be_read_gets_an_honest_answer(ctx): | |
| await say(ctx, OWNER, "/resize 320") | |
| ev = await send_photo(ctx, data=b"not a picture at all") | |
| assert "couldn't read that image" in ev.texts[0] and results(ctx) == [] | |
| ctx.client.media[50] = None | |
| await say(ctx, OWNER, "/resize 320") | |
| ev = await send_photo(ctx, data=b"") | |
| assert "couldn't download" in ev.texts[0] | |
| async def test_very_large_files_are_refused_before_downloading(ctx): | |
| await say(ctx, OWNER, "/resize 320") | |
| ev = await send_image_file(ctx, size=appmod.RESIZE_MAX_BYTES + 1) | |
| assert "I can resize images up to 20.00 MB" in ev.texts[0] and ctx.client.downloads == [] | |
| async def test_it_expires_and_cancel_ends_it(ctx): | |
| await say(ctx, OWNER, "/resize 320") | |
| ctx.app.modes[OWNER].expires = time.time() - 1 | |
| ev = await send_doc(ctx, OWNER, "A.mkv", mid=60) | |
| assert "Queued" in ev.texts[0] | |
| await say(ctx, OWNER, "/resize 320") | |
| ev = await say(ctx, OWNER, "/cancel") | |
| assert "Cancelled" in ev.texts[0] and OWNER not in ctx.app.modes | |
| await send_photo(ctx) | |
| assert results(ctx) == [] and await auto_thumb(ctx) | |
| async def test_a_format_thumbnail_prompt_takes_priority_over_a_waiting_resize(ctx): | |
| fmt = await ctx.db.create_format(OWNER, "One Piece") | |
| await say(ctx, OWNER, "/resize 320") | |
| await click(ctx, OWNER, f"fm:th:{fmt['id']}") # opened later: the user is setting the format's thumbnail now | |
| await dispatch(ctx.client, FakeEvent(OWNER, message=FakeMessage("", photo=True, mid=70), photo_bytes=jpeg())) | |
| assert (await ctx.db.get_format(OWNER, fmt["id"]))["has_thumb"] and results(ctx) == [] | |
| async def test_strangers_cannot_use_it(ctx): | |
| ev = await say(ctx, STRANGER, "/resize 320") | |
| 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, "/resize 320") | |
| assert "1 file(s) was discarded" in ev.texts[0] and ctx.app.modes[OWNER].kind == "resize" | |
| cb = await click(ctx, OWNER, "rz:s:640") # via the menu too | |
| assert "discarded" not in cb.edits[-1][0] # the sequence was already dropped by the first command | |
| async def test_both_commands_are_in_the_command_menu(ctx): | |
| assert "resize" in dict(appmod.COMMANDS) and "ethumb" in dict(appmod.COMMANDS) | |