"""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"] == "[S01E05] Tokyo Revengers [1080p]" 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="THE AUTO FORMAT'S {filename}") # 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="{filename} · {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"One Piece - E5 · {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="{filename} · {duration}")) assert client.sent[0][2]["caption"] == f"[S01E05] Tokyo Revengers [1080p] · {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="{filename} · {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"[S01E05] Tokyo Revengers [1080p] · {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"] == "[S01E05] Tokyo Revengers [1080p]" 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("{filename} · {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"New - E5 · [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("{2episode}", new_name="x.mkv", orig_name="Plain.mkv", ext=".mkv", size=1, duration=0) == "" 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