Blame for tests/test_stream_rename.py

"""Streaming renames: a big file goes straight from the download into the upload, and anything that needs the file takes the usual way."""
import shutil
from pathlib import Path
from types import SimpleNamespace

import pytest
from telethon import errors, types

from bot import errorlog
from bot import fasttransfer as ft
from bot import pipeline
from bot.config import load_config
from bot.db import Database
from bot.errors import Cancelled
from bot.media import probe
from bot.pipeline import Renamer, embedded_thumb, make_caption, source_attributes
from bot.queue import Job
from bot.utils import human_duration, human_size

from .appkit import ensure_auto
from .conftest import needs_ffmpeg

MIB = 1024 * 1024
BIG = 40 * MIB  # more than the 10 MB Telegram's big-file upload starts at; nothing this size is really moved: the engine is a stub here
NAME = "[SubsPlease] Tokyo Revengers - 05 (1080p) [ABCD1234].mkv"
NEW_NAME = "[S01E05] Tokyo Revengers [1080p].mkv"
PREVIEW = b"\xff\xd8preview-of-the-original"


def video_attr(duration=1400.5, w=1920, h=1080, **kw):
    return types.DocumentAttributeVideo(duration=duration, w=w, h=h, supports_streaming=True, **kw)


def audio_attr(duration=215.4, title="Song", performer="Me", **kw):
    return types.DocumentAttributeAudio(duration=duration, title=title, performer=performer, **kw)


def preview_sizes():
    return [types.PhotoStrippedSize(type="i", bytes=b"\x01\x02"), types.PhotoSize(type="m", w=320, h=180, size=900)]


class FakeMessage:
    """A message with a document, the way Telethon shows it to the pipeline."""

    def __init__(self, name=NAME, size=BIG, *, attributes=(), thumbs=None, mime="video/x-matroska", msg_id=7, document=True):
        self.id = msg_id
        self.file = SimpleNamespace(name=name, size=size, ext=Path(name or "").suffix, mime_type=mime)
        self.document = SimpleNamespace(attributes=list(attributes), thumbs=thumbs, size=size, mime_type=mime) if document else None
        video, audio = source_attributes(self)
        self.video = self.document if video is not None else None  # like Telethon's own properties
        self.audio = self.document if audio is not None else None


class FakeClient:
    def __init__(self, message, source):
        self.message, self.source = message, source
        self.sent, self.edits = [], []
        self.relays, self.downloads, self.uploads, self.preview_downloads = [], [], [], 0
        self.relay_error = None
        self.send_errors = []

    async def get_messages(self, chat_id, ids):
        return self.message

    async def download_media(self, message, file=None, thumb=None, progress_callback=None):
        assert thumb is not None, "the whole file is fetched by transfer.download, which is stubbed"
        self.preview_downloads += 1
        return PREVIEW

    async def send_file(self, target, input_file, **kwargs):
        if self.send_errors:
            raise self.send_errors.pop(0)
        self.sent.append((target, input_file, kwargs))

    async def edit_message(self, message, text, buttons=None, parse_mode=None):
        self.edits.append((text, buttons))


@pytest.fixture(autouse=True)
def fresh_stats(monkeypatch):
    monkeypatch.setattr(ft, "stats", ft.TransferStats())


@pytest.fixture
async def env(tmp_path, media_dir, monkeypatch):
    def new_cfg(**environ):
        return load_config({
            "API_ID": "1", "API_HASH": "h", "BOT_TOKEN": "1:t", "OWNER_ID": "42", "TEMP_DIR": str(tmp_path / "work"),
            "PROGRESS_INTERVAL": "2", **environ,  # fast transfer and streaming are on unless a test turns them off
        })

    db = Database(f"sqlite+aiosqlite:///{tmp_path}/s.db")
    await db.init()
    ns = SimpleNamespace(db=db, media_dir=media_dir, tmp=tmp_path, new_cfg=new_cfg, cfg=new_cfg())

    async def relay(client, document, file_name, *, connections, progress=None, cancel=None):
        client.relays.append((document, file_name, connections))
        if client.relay_error is not None:
            raise client.relay_error
        progress(document.size // 2, document.size)
        progress(document.size, document.size)
        return types.InputFileBig(id=99, parts=-(-document.size // ft.UPLOAD_PART), name=file_name)

    async def download(client, message, dest, *, fast, connections, progress=None, cancel=None, max_bytes=None):
        client.downloads.append(Path(dest).name)
        shutil.copy(client.source, dest)

    async def upload(client, path, file_name, *, fast, connections, progress=None, cancel=None):
        client.uploads.append(Path(path).read_bytes())
        return types.InputFileBig(id=1, parts=1, name=file_name)

    monkeypatch.setattr(ft, "relay", relay)
    monkeypatch.setattr(ft, "download", download)
    monkeypatch.setattr(ft, "upload", upload)
    yield ns
    await db.close()


@pytest.fixture
async def log_channel():
    """What reaches the log channel while a test runs."""
    reports = []

    async def send(text):
        reports.append(text)

    handler = errorlog.install(send)
    yield SimpleNamespace(reports=reports, flush=handler.aclose)
    handler.detach()


async def rename(env, message=None, *, cfg=None, new_name=NEW_NAME, fmt=None, source="sample.mkv"):
    """Run one job; returns ``(client, job)``."""
    message = message or FakeMessage(attributes=[video_attr()], thumbs=preview_sizes())
    client = FakeClient(message, env.media_dir / source)
    job = Job(id=1, user_id=42, chat_id=42, msg_id=7, file_name=message.file.name or "x", size=message.file.size, new_name=new_name,
              status_msg=SimpleNamespace(id=99), fmt=fmt)
    await Renamer(client, cfg or env.cfg, env.db).process(job)
    return client, job


def attributes_by_type(kwargs):
    return {type(a): a for a in kwargs["attributes"]}


async def auto(env, thumb=None, **changes):
    """The user's auto format (made on demand) with these changes and this thumbnail: what renames the files they just send."""
    fmt = await ensure_auto(env.db, 42)
    if changes:
        await env.db.update_format(42, fmt["id"], **changes)
    if thumb:
        await env.db.set_format_thumb(42, fmt["id"], thumb)
    return await env.db.get_format(42, fmt["id"])


# ---------------------------------------------------------------------------
# the streamed way
# ---------------------------------------------------------------------------


async def test_a_big_plain_rename_is_streamed_with_the_details_telegram_already_has(env):
    message = FakeMessage(attributes=[video_attr()], thumbs=preview_sizes())
    client, job = await rename(env, message)

    assert client.relays == [(message.document, NEW_NAME, env.cfg.connections)]
    assert client.downloads == [] and client.uploads == []  # nothing was saved, and nothing uploaded from a file
    (target, input_file, kw), = client.sent
    assert target == 42 and kw["reply_to"] == 7 and input_file.name == NEW_NAME
    by_type = attributes_by_type(kw)
    assert by_type[types.DocumentAttributeFilename].file_name == NEW_NAME
    video = by_type[types.DocumentAttributeVideo]
    assert (video.duration, video.w, video.h, video.supports_streaming) == (1400.5, 1920, 1080, True)
    assert kw["force_document"] is False and kw["mime_type"] == "video/x-matroska" and kw["parse_mode"] == "html"
    assert kw["thumb"] == PREVIEW and client.preview_downloads == 1  # the original's own preview, as no frame can be cut from a file
    assert kw["caption"] == "<b>[S01E05] Tokyo Revengers [1080p]</b>"
    assert job.state == "done" and (await env.db.totals())[0] == 1
    assert "Streaming" in client.edits[0][0]  # the progress message says what is going on
    done, buttons = client.edits[-1]
    assert "Done" in done and human_size(BIG) in done and buttons is None
    assert ft.stats.streamed == 1 and ft.stats.stream_fallbacks == 0
    assert not (Path(env.cfg.temp_dir) / "job1").exists()


async def test_the_formats_thumbnail_wins_and_the_original_preview_is_then_not_even_fetched(env):
    client, _job = await rename(env, fmt=await auto(env, thumb=b"\xff\xd8mine"))
    assert client.sent[0][2]["thumb"] == b"\xff\xd8mine" and client.preview_downloads == 0 and len(client.relays) == 1

    fmt = await env.db.create_format(42, "One Piece")
    await env.db.set_format_thumb(42, fmt["id"], b"\xff\xd8format")
    client, _job = await rename(env, fmt=await env.db.get_format(42, fmt["id"]))
    assert client.sent[0][2]["thumb"] == b"\xff\xd8format" and len(client.relays) == 1


async def test_a_format_streams_with_its_own_caption(env):
    await auto(env, caption="<i>THE AUTO FORMAT'S {filename}</i>")  # another format's caption is not borrowed
    fmt = await env.db.create_format(42, "One Piece", template="One Piece - E{episode}")
    await env.db.update_format(42, fmt["id"], caption="<b>{filename}</b> · {size} · {duration}")
    client, _job = await rename(env, new_name="One Piece - E5.mkv", fmt=await env.db.get_format(42, fmt["id"]))
    (_t, _f, kw), = client.sent
    assert len(client.relays) == 1
    assert kw["caption"] == f"<b>One Piece - E5</b> · {human_size(BIG)} · {human_duration(1400.5)}"  # the duration is the one Telegram has


async def test_a_file_that_gets_no_thumbnail_anyway_is_streamed_without_one(env):
    message = FakeMessage("Archive.zip", BIG, mime="application/zip")
    client, _job = await rename(env, message, new_name="Archive 2.zip")
    (_t, _f, kw), = client.sent
    assert len(client.relays) == 1 and kw["thumb"] is None and client.preview_downloads == 0
    assert kw["force_document"] is True and kw["mime_type"] == "application/zip"
    assert [type(a) for a in kw["attributes"]] == [types.DocumentAttributeFilename]


async def test_an_original_sent_as_a_file_is_streamed_as_a_file(env):
    client, _job = await rename(env, FakeMessage(), fmt=await auto(env, thumb=b"\xff\xd8mine"))  # an .mkv without any video details: Telegram sent it as a file
    (_t, _f, kw), = client.sent
    assert len(client.relays) == 1 and kw["force_document"] is True and kw["thumb"] == b"\xff\xd8mine"
    assert types.DocumentAttributeVideo not in attributes_by_type(kw)


async def test_audio_is_streamed_with_the_title_and_artist_telegram_has(env):
    message = FakeMessage("old name.mp3", BIG, attributes=[audio_attr()], mime="audio/mpeg")
    client, _job = await rename(env, message, new_name="New Song.mp3")
    (_t, _f, kw), = client.sent
    audio = attributes_by_type(kw)[types.DocumentAttributeAudio]
    assert (audio.duration, audio.title, audio.performer) == (215, "Song", "Me") and not audio.voice
    assert kw["force_document"] is False and kw["mime_type"] == "audio/mpeg" and kw["thumb"] is None and len(client.relays) == 1


async def test_a_forced_type_is_followed_when_telegram_has_the_details_for_it(env):
    await env.db.update_settings(42, file_type="document")
    client, _job = await rename(env)  # sent as a video, but the user wants documents
    (_t, _f, kw), = client.sent
    assert len(client.relays) == 1 and kw["force_document"] is True and types.DocumentAttributeVideo not in attributes_by_type(kw)


async def test_the_caption_may_use_the_duration_when_telegram_has_it(env):
    client, _job = await rename(env, fmt=await auto(env, caption="<b>{filename}</b> · {duration}"))
    assert client.sent[0][2]["caption"] == f"<b>[S01E05] Tokyo Revengers [1080p]</b> · {human_duration(1400.5)}" and len(client.relays) == 1


async def test_a_streamed_file_goes_to_the_target_channel_like_any_other(env):
    await env.db.update_settings(42, target_channel=-100123)
    client, _job = await rename(env)
    (target, _f, kw), = client.sent
    assert target == -100123 and "reply_to" not in kw and "channel" in client.edits[-1][0]


async def test_a_rejected_caption_is_dropped_for_a_streamed_file_too(env):
    message = FakeMessage(attributes=[video_attr()], thumbs=preview_sizes())
    client = FakeClient(message, env.media_dir / "sample.mkv")
    client.send_errors = [errors.MediaCaptionTooLongError(request=None)]
    job = Job(id=1, user_id=42, chat_id=42, msg_id=7, file_name=NAME, size=BIG, new_name=NEW_NAME, status_msg=SimpleNamespace(id=99))
    await Renamer(client, env.cfg, env.db).process(job)
    (_t, _f, kw), = client.sent
    assert kw["caption"] is None and "caption was rejected" in client.edits[-1][0] and job.state == "done" and len(client.relays) == 1


async def test_the_done_message_gives_the_speed_of_the_whole_stream(env, monkeypatch):
    ticks = iter(range(0, 1000, 2))  # every look at the clock is two seconds later than the one before
    monkeypatch.setattr(pipeline, "time", SimpleNamespace(monotonic=lambda: next(ticks)))
    client, _job = await rename(env)
    assert f"{human_size(BIG)} · ⬇⬆ {human_size(BIG / 2)}/s streamed" in client.edits[-1][0]


async def test_a_stream_that_took_under_a_second_shows_no_speed(env):
    client, _job = await rename(env)
    assert "Done" in client.edits[-1][0] and "/s" not in client.edits[-1][0]


# ---------------------------------------------------------------------------
# when the file has to be looked at: the usual way (download, ffmpeg, upload)
# ---------------------------------------------------------------------------


def usual_way(client, job):
    """The file was saved and uploaded the usual way, and sent once."""
    assert client.relays == [] and len(client.downloads) == 1 and len(client.uploads) == 1
    assert len(client.sent) == 1 and job.state == "done"
    assert ft.stats.streamed == 0 and ft.stats.stream_fallbacks == 0  # not streaming is not a failure


@needs_ffmpeg
async def test_small_files_are_not_streamed(env):
    message = FakeMessage(size=ft.BIG_FILE_THRESHOLD, attributes=[video_attr()], thumbs=preview_sizes())
    client, job = await rename(env, message)
    usual_way(client, job)


@needs_ffmpeg
async def test_metadata_and_tracks_change_the_file_so_it_is_not_streamed(env):
    await env.db.update_settings(42, metadata_on=True, meta={"title": "Tokyo Revengers"})
    client, job = await rename(env)
    usual_way(client, job)


@needs_ffmpeg
async def test_streaming_can_be_switched_off_and_needs_the_parallel_transfer(env):
    for environ in ({"STREAM_RENAME": "0"}, {"FAST_TRANSFER": "0"}):
        client, job = await rename(env, cfg=env.new_cfg(**environ))
        usual_way(client, job)


@needs_ffmpeg
async def test_a_video_without_a_thumbnail_to_reuse_is_saved_so_that_one_can_be_cut_from_it(env):
    for thumbs in (None, [types.PhotoStrippedSize(type="i", bytes=b"\x01\x02")]):  # a blurred stripped preview is no thumbnail
        message = FakeMessage(attributes=[video_attr()], thumbs=thumbs)
        client, job = await rename(env, message)
        usual_way(client, job)
        assert client.sent[0][2]["thumb"][:2] == b"\xff\xd8" and client.sent[0][2]["thumb"] != PREVIEW  # a frame of the video


@needs_ffmpeg
async def test_a_forced_video_or_audio_type_needs_the_details_only_the_file_can_give(env):
    await env.db.update_settings(42, file_type="video")
    client, job = await rename(env, FakeMessage(thumbs=preview_sizes()))  # an .mkv sent as a file: no video details
    usual_way(client, job)
    assert types.DocumentAttributeVideo in attributes_by_type(client.sent[0][2])  # probed from the file instead

    await env.db.update_settings(42, file_type="audio")
    client, job = await rename(env, FakeMessage(attributes=[video_attr()], thumbs=preview_sizes()))
    usual_way(client, job)


@needs_ffmpeg
async def test_a_caption_that_wants_the_duration_needs_the_file_when_telegram_does_not_know_it(env):
    fmt = await auto(env, thumb=b"\xff\xd8mine", caption="<b>{filename}</b> · {duration}")
    client, job = await rename(env, FakeMessage(), fmt=fmt)  # sent as a file: no duration
    usual_way(client, job)
    duration = (await probe(env.media_dir / "sample.mkv")).duration
    assert client.sent[0][2]["caption"] == f"<b>[S01E05] Tokyo Revengers [1080p]</b> · {human_duration(duration)}"  # read from the file


@needs_ffmpeg
async def test_a_message_without_a_document_is_never_streamed(env):
    client, job = await rename(env, FakeMessage(document=False))
    usual_way(client, job)


# ---------------------------------------------------------------------------
# when streaming goes wrong: the usual way takes over
# ---------------------------------------------------------------------------


@needs_ffmpeg
async def test_a_failed_stream_falls_back_to_the_usual_way_and_says_so(env, log_channel):
    client = FakeClient(FakeMessage(attributes=[video_attr()], thumbs=preview_sizes()), env.media_dir / "sample.mkv")
    client.relay_error = RuntimeError("DC unreachable")
    job = Job(id=1, user_id=42, chat_id=42, msg_id=7, file_name=NAME, size=BIG, new_name=NEW_NAME, status_msg=SimpleNamespace(id=99))
    await Renamer(client, env.cfg, env.db).process(job)
    await log_channel.flush()

    assert len(client.relays) == 1 and len(client.downloads) == 1 and client.uploads == [(env.media_dir / "sample.mkv").read_bytes()]
    (_t, _f, kw), = client.sent
    assert kw["force_document"] is False and kw["caption"] == "<b>[S01E05] Tokyo Revengers [1080p]</b>"
    assert job.state == "done" and (await env.db.totals())[0] == 1  # counted once
    done = client.edits[-1][0]
    assert "Done" in done and "Streaming didn't work for this file" in done
    assert (ft.stats.streamed, ft.stats.stream_fallbacks, ft.stats.last_stream_fallback) == (0, 1, "RuntimeError: DC unreachable")
    assert ft.stats.fallbacks == 0  # the parallel transfer itself did not fail
    (report,) = log_channel.reports  # the log channel is told, with the traceback
    assert "streaming" in report and "RuntimeError: DC unreachable" in report and "Traceback" in report
    assert not (Path(env.cfg.temp_dir) / "job1").exists()


@needs_ffmpeg
async def test_telegram_refusing_the_streamed_parts_also_falls_back(env):
    client = FakeClient(FakeMessage(attributes=[video_attr()], thumbs=preview_sizes()), env.media_dir / "sample.mkv")
    client.send_errors = [errors.FilePartMissingError(request=None, capture=3)]
    job = Job(id=1, user_id=42, chat_id=42, msg_id=7, file_name=NAME, size=BIG, new_name=NEW_NAME, status_msg=SimpleNamespace(id=99))
    await Renamer(client, env.cfg, env.db).process(job)
    assert len(client.relays) == 1 and len(client.downloads) == 1 and len(client.uploads) == 1 and len(client.sent) == 1
    assert job.state == "done" and "Streaming didn't work" in client.edits[-1][0] and (await env.db.totals())[0] == 1
    assert (ft.stats.streamed, ft.stats.stream_fallbacks) == (0, 1) and "FilePartMissingError" in ft.stats.last_stream_fallback


async def test_any_other_refusal_after_streaming_fails_the_job_without_a_second_try(env):
    client = FakeClient(FakeMessage(attributes=[video_attr()], thumbs=preview_sizes()), env.media_dir / "sample.mkv")
    client.send_errors = [errors.ChatWriteForbiddenError(request=None)]
    job = Job(id=1, user_id=42, chat_id=42, msg_id=7, file_name=NAME, size=BIG, new_name=NEW_NAME, status_msg=SimpleNamespace(id=99))
    await Renamer(client, env.cfg, env.db).process(job)
    assert job.state == "failed" and client.sent == [] and client.downloads == [] and len(client.relays) == 1  # no download/upload all over again
    assert "Failed" in client.edits[-1][0] and ft.stats.stream_fallbacks == 0 and (await env.db.totals())[0] == 0


async def test_cancelling_a_stream_cancels_the_job_and_does_not_fall_back(env):
    client = FakeClient(FakeMessage(attributes=[video_attr()], thumbs=preview_sizes()), env.media_dir / "sample.mkv")
    client.relay_error = Cancelled()
    job = Job(id=1, user_id=42, chat_id=42, msg_id=7, file_name=NAME, size=BIG, new_name=NEW_NAME, status_msg=SimpleNamespace(id=99))
    await Renamer(client, env.cfg, env.db).process(job)
    assert job.state == "cancelled" and client.sent == [] and client.downloads == [] and "Cancelled" in client.edits[-1][0]
    assert ft.stats.stream_fallbacks == 0 and ft.stats.streamed == 0


async def test_a_cancel_that_arrives_as_the_stream_ends_sends_nothing(env, monkeypatch):
    client = FakeClient(FakeMessage(attributes=[video_attr()], thumbs=preview_sizes()), env.media_dir / "sample.mkv")
    job = Job(id=1, user_id=42, chat_id=42, msg_id=7, file_name=NAME, size=BIG, new_name=NEW_NAME, status_msg=SimpleNamespace(id=99))
    real_relay = ft.relay

    async def relay_then_cancel(*a, **kw):
        result = await real_relay(*a, **kw)
        job.cancel.set()
        return result

    monkeypatch.setattr(ft, "relay", relay_then_cancel)
    await Renamer(client, env.cfg, env.db).process(job)
    assert job.state == "cancelled" and client.sent == []


# ---------------------------------------------------------------------------
# helpers
# ---------------------------------------------------------------------------


def test_source_attributes_ignore_round_videos_and_voice_notes():
    plain, round_video = video_attr(), video_attr(round_message=True)
    song, voice = audio_attr(), audio_attr(voice=True)
    assert source_attributes(FakeMessage(attributes=[plain, song])) == (plain, song)
    assert source_attributes(FakeMessage(attributes=[round_video, voice])) == (None, None)
    assert source_attributes(FakeMessage(attributes=[types.DocumentAttributeFilename("a.mkv")])) == (None, None)
    assert source_attributes(FakeMessage(document=False)) == (None, None) and source_attributes(object()) == (None, None)


def test_make_caption_fills_in_what_it_is_given():
    text = make_caption("<b>{filename}</b> · {original} · {size} · {duration} · {ext} · {nonsense}", new_name="New - E5.mkv",
                        orig_name="[Grp] Show - 05 [720p].mkv", ext=".mkv", size=3 * MIB, duration=3725)
    assert text == f"<b>New - E5</b> · [Grp] Show - 05 [720p].mkv · {human_size(3 * MIB)} · {human_duration(3725)} · mkv · {{nonsense}}"
    assert make_caption("{duration}", new_name="a.mkv", orig_name="a.mkv", ext=".mkv", size=1, duration=0) is None  # empty: no caption
    assert len(make_caption("x" * 5000, new_name="a.mkv", orig_name="a.mkv", ext=".mkv", size=1, duration=0)) == 1000


def test_make_caption_pads_the_episode_when_asked_to():
    text = make_caption("Ep {episode} / {1episode} / {2episode} / {3episode}", new_name="x.mkv",
                        orig_name="Show - 7 [720p].mkv", ext=".mkv", size=1, duration=0)
    assert text == "Ep 07 / 07 / 007 / 0007"
    assert make_caption("{3episode} {3episode}", new_name="x.mkv", orig_name="Show - 321 [720p].mkv", ext=".mkv", size=1, duration=0) == "0321 0321"
    assert make_caption("{3episode}", new_name="x.mkv", orig_name="Show - 1321 [720p].mkv", ext=".mkv", size=1, duration=0) == "1321"
    assert make_caption("<b>{2episode}</b>", new_name="x.mkv", orig_name="Plain.mkv", ext=".mkv", size=1, duration=0) == "<b></b>"


async def test_the_embedded_preview_is_only_fetched_when_it_is_worth_having():
    class Client:
        def __init__(self, data=PREVIEW, error=None):
            self.data, self.error, self.calls = data, error, 0

        async def download_media(self, msg, file=None, thumb=None):
            self.calls += 1
            if self.error:
                raise self.error
            return self.data

    stripped = [types.PhotoStrippedSize(type="i", bytes=b"\x01\x02")]
    client = Client()
    assert await embedded_thumb(client, FakeMessage(thumbs=preview_sizes())) == PREVIEW
    assert await embedded_thumb(client, FakeMessage(thumbs=stripped)) is None and client.calls == 1  # blurred: not asked for
    assert await embedded_thumb(client, FakeMessage(thumbs=stripped), allow_stripped=True) == PREVIEW
    assert await embedded_thumb(client, FakeMessage(thumbs=None)) is None and await embedded_thumb(client, object()) is None
    assert await embedded_thumb(Client(error=OSError("gone")), FakeMessage(thumbs=preview_sizes())) is None
    assert await embedded_thumb(Client(data=b""), FakeMessage(thumbs=preview_sizes())) is None