Blame for tests/test_resize.py

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