import shutil from pathlib import Path from types import SimpleNamespace import pytest from telethon import errors, types from bot import errorlog from bot.config import load_config from bot.db import Database from bot.media import probe from bot.parser import DEFAULT_TEMPLATE, compute_new_name from bot.pipeline import Renamer, decide_send_type, guess_mime, original_name_of, source_ext from bot.queue import Job from .appkit import ensure_auto from .conftest import needs_ffmpeg class FakeMessage: def __init__(self, name, size, *, video=True, audio=False, mime="video/x-matroska", msg_id=7, ext=None): self.id = msg_id self.file = SimpleNamespace(name=name, size=size, ext=ext if ext is not None else Path(name or "").suffix, mime_type=mime) self.video = object() if video else None self.audio = object() if audio else None self.document = object() class FakeClient: def __init__(self, source: Path, message: FakeMessage): self.source = source self.message = message self.sent = [] self.edits = [] self.uploads = {} self.download_error = None self.send_errors = [] async def get_messages(self, chat_id, ids): return self.message async def download_media(self, message, file, progress_callback): progress_callback(1, 2) if self.download_error: raise self.download_error shutil.copy(self.source, file) progress_callback(2, 2) async def upload_file(self, path, file_name, progress_callback): progress_callback(1, 1) keep = Path(path).with_name("uploaded_copy_" + Path(path).name) keep = Path(self.source).parent / ("upload-" + str(len(self.uploads)) + Path(path).suffix) shutil.copy(path, keep) self.uploads[len(self.uploads)] = keep return types.InputFileBig(id=len(self.uploads), parts=1, name=file_name) 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 async def env(tmp_path, media_dir): cfg = load_config({ "API_ID": "1", "API_HASH": "h", "BOT_TOKEN": "1:t", "OWNER_ID": "42", "TEMP_DIR": str(tmp_path / "work"), "FAST_TRANSFER": "0", "PROGRESS_INTERVAL": "2", }) db = Database(f"sqlite+aiosqlite:///{tmp_path}/p.db") await db.init() ns = SimpleNamespace(cfg=cfg, db=db, media_dir=media_dir, tmp=tmp_path) yield ns await db.close() def make(env, name="[SubsPlease] Tokyo Revengers - 05 (1080p) [ABCD1234].mkv", src="sample.mkv", **msg_kwargs): source = env.media_dir / src message = FakeMessage(name, source.stat().st_size, **msg_kwargs) client = FakeClient(source, message) renamer = Renamer(client, env.cfg, env.db) new_name, _ = compute_new_name(template=DEFAULT_TEMPLATE, mode="filename", original_name=name, original_ext=Path(name).suffix) job = Job(id=1, user_id=42, chat_id=42, msg_id=7, file_name=name, size=source.stat().st_size, new_name=new_name, status_msg=SimpleNamespace(id=99)) return client, renamer, job @needs_ffmpeg async def test_plain_rename_sends_video_with_thumbnail(env): client, renamer, job = make(env) await renamer.process(job) assert job.new_name == "[S01E05] Tokyo Revengers [1080p].mkv" (target, input_file, kw), = client.sent assert target == 42 and kw["reply_to"] == 7 by_type = {type(a): a for a in kw["attributes"]} assert by_type[types.DocumentAttributeFilename].file_name == job.new_name video = by_type[types.DocumentAttributeVideo] assert (video.w, video.h) == (320, 240) and 1 < video.duration < 3 and video.supports_streaming assert kw["force_document"] is False and kw["mime_type"] == "video/x-matroska" assert kw["thumb"][:2] == b"\xff\xd8" # generated from a video frame assert kw["caption"] == "[S01E05] Tokyo Revengers [1080p]" and kw["parse_mode"] == "html" # bold, by default assert input_file.name == job.new_name assert client.uploads[0].read_bytes() == (env.media_dir / "sample.mkv").read_bytes() # untouched without metadata assert "Done" in client.edits[-1][0] and client.edits[-1][1] is None assert job.state == "done" assert (await env.db.totals())[0] == 1 assert not (Path(env.cfg.temp_dir) / "job1").exists() # temp files cleaned @needs_ffmpeg async def test_notes_from_queueing_appear_in_the_done_message(env): client, renamer, job = make(env) job.notes = ["Couldn't find {episode} in the file name, so it is left empty."] await renamer.process(job) assert job.state == "done" assert "⚠️ Couldn't find {episode} in the file name" in client.edits[-1][0] @needs_ffmpeg async def test_metadata_and_track_selection_are_applied_before_upload(env): await env.db.update_settings( 42, metadata_on=True, keep_audio=["english", "japanese"], keep_subtitle=["2"], meta={"title": "Tokyo Revengers", "audio": "{lang}", "subtitle": "English"}, ) client, renamer, job = make(env) await renamer.process(job) out = await probe(client.uploads[0]) assert [s.lang for s in out.audio] == ["eng", "jpn"] assert [s.lang for s in out.subtitles] == ["eng"] assert out.tags.get("title") == "Tokyo Revengers" assert [s.title for s in out.audio] == ["English", "Japanese"] assert client.sent and job.state == "done" @needs_ffmpeg async def test_document_file_type_is_forced(env): await env.db.update_settings(42, file_type="document") client, renamer, job = make(env) await renamer.process(job) (_t, _f, kw), = client.sent assert kw["force_document"] is True assert types.DocumentAttributeVideo not in {type(a) for a in kw["attributes"]} @needs_ffmpeg async def test_same_type_keeps_document_as_document(env): client, renamer, job = make(env, video=False) # original was sent as a file await renamer.process(job) (_t, _f, kw), = client.sent assert kw["force_document"] is True assert types.DocumentAttributeVideo not in {type(a) for a in kw["attributes"]} @needs_ffmpeg async def test_forced_video_for_non_video_falls_back_to_document(env, tmp_path): blob = tmp_path / "data.zip" blob.write_bytes(b"PK" + b"0" * 5000) message = FakeMessage("stuff.zip", 5002, video=False, mime="application/zip") client = FakeClient(blob, message) renamer = Renamer(client, env.cfg, env.db) await env.db.update_settings(42, file_type="video") job = Job(id=2, user_id=42, chat_id=42, msg_id=7, file_name="stuff.zip", size=5002, new_name="Stuff.zip", status_msg=object()) await renamer.process(job) (_t, _f, kw), = client.sent assert kw["force_document"] is True and kw["mime_type"] == "application/zip" assert kw["thumb"] is None async def make_auto(env, thumb=None, **changes): """The user's auto format (made on demand), with these changes and this thumbnail.""" 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"]) @needs_ffmpeg async def test_custom_thumbnail_wins_over_generated(env): client, renamer, job = make(env) job.fmt = await make_auto(env, thumb=b"\xff\xd8custom") await renamer.process(job) assert client.sent[0][2]["thumb"] == b"\xff\xd8custom" @needs_ffmpeg async def test_target_channel(env): await env.db.update_settings(42, target_channel=-100123) client, renamer, job = make(env) await renamer.process(job) (target, _f, kw), = client.sent assert target == -100123 and "reply_to" not in kw assert "channel" in client.edits[-1][0] @needs_ffmpeg async def test_audio_file_sent_as_audio_with_tags(env): await env.db.update_settings(42, metadata_on=True, meta={"title": "Song", "artist": "Me"}) client, renamer, job = make(env, name="old name.mp3", src="sample.mp3", video=False, audio=True, mime="audio/mpeg") job.new_name = "New Song.mp3" await renamer.process(job) (_t, _f, kw), = client.sent audio = next(a for a in kw["attributes"] if isinstance(a, types.DocumentAttributeAudio)) assert audio.title == "Song" and audio.performer == "Me" and audio.duration >= 1 assert kw["force_document"] is False and kw["mime_type"] == "audio/mpeg" assert (await probe(client.uploads[0])).tags["title"] == "Song" @needs_ffmpeg async def test_unsupported_container_with_metadata_warns_but_still_renames(env, tmp_path): # an .avi has video streams but is outside the supported remux containers avi = tmp_path / "x.avi" import subprocess subprocess.run(["ffmpeg", "-y", "-loglevel", "error", "-f", "lavfi", "-i", "testsrc=duration=1:size=64x64:rate=5", "-c:v", "mpeg4", str(avi)], check=True) message = FakeMessage("x.avi", avi.stat().st_size, video=False, mime="video/x-msvideo") client = FakeClient(avi, message) await env.db.update_settings(42, metadata_on=True, meta={"title": "T"}) job = Job(id=3, user_id=42, chat_id=42, msg_id=7, file_name="x.avi", size=avi.stat().st_size, new_name="Y.avi", status_msg=object()) await Renamer(client, env.cfg, env.db).process(job) assert client.sent and "not supported" in client.edits[-1][0] @needs_ffmpeg async def test_caption_rejected_is_retried_without_caption(env): client, renamer, job = make(env) client.send_errors = [errors.MediaCaptionTooLongError(request=None)] await renamer.process(job) (_t, _f, kw), = client.sent assert kw["caption"] is None assert "caption was rejected" in client.edits[-1][0] and job.state == "done" @needs_ffmpeg async def test_cancel_before_start(env): client, renamer, job = make(env) job.cancel.set() await renamer.process(job) assert job.state == "cancelled" and client.sent == [] assert "Cancelled" in client.edits[-1][0] assert not (Path(env.cfg.temp_dir) / "job1").exists() @pytest.fixture async def log_channel(): """What reaches the log channel while a test runs (the error reporter is attached to the root logger, as in main.py).""" reports = [] async def send(text): reports.append(text) handler = errorlog.install(send) yield SimpleNamespace(reports=reports, flush=handler.aclose) handler.detach() @needs_ffmpeg async def test_failure_is_reported_and_logged(env, log_channel): client, renamer, job = make(env) client.download_error = RuntimeError("network down") await renamer.process(job) await log_channel.flush() assert job.state == "failed" and client.sent == [] assert "Failed" in client.edits[-1][0] and "network down" in client.edits[-1][0] (report,) = log_channel.reports # the log channel gets the whole traceback, as a Python code block assert report.startswith(f'
Error processing file:\nuser 42 · job #1\n{job.file_name} → {job.new_name} (')
    assert "Traceback (most recent call last):" in report and report.endswith("RuntimeError: network down
") assert not (Path(env.cfg.temp_dir) / "job1").exists() @needs_ffmpeg async def test_a_finished_job_leaves_nothing_in_the_log_channel(env, log_channel): client, renamer, job = make(env) await renamer.process(job) await log_channel.flush() assert job.state == "done" and log_channel.reports == [] @needs_ffmpeg async def test_not_enough_disk_space(env, monkeypatch): client, renamer, job = make(env) monkeypatch.setattr(shutil, "disk_usage", lambda p: SimpleNamespace(total=100, used=100, free=10)) await renamer.process(job) assert job.state == "failed" and "disk space" in client.edits[-1][0] @needs_ffmpeg async def test_deleted_original_message_fails_cleanly(env): client, renamer, job = make(env) async def gone(chat_id, ids): return None client.get_messages = gone await renamer.process(job) assert job.state == "failed" and "deleted" in client.edits[-1][0] def test_helpers(): video = FakeMessage("a.mkv", 1) doc = FakeMessage("a.mkv", 1, video=False) audio = FakeMessage("a.mp3", 1, video=False, audio=True) assert decide_send_type("same", video) == "video" assert decide_send_type("same", doc) == "document" assert decide_send_type("same", audio) == "audio" assert decide_send_type("document", video) == "document" assert decide_send_type("audio", doc) == "audio" assert guess_mime("x.mkv", None) == "video/x-matroska" assert guess_mime("x.unknown123", "foo/bar") == "foo/bar" assert guess_mime("x.unknown123", None) == "application/octet-stream" nameless = FakeMessage(None, 1, ext=".mp4", msg_id=55) assert original_name_of(nameless) == "video_55.mp4" assert source_ext(nameless, "video_55.mp4") == ".mp4" assert source_ext(FakeMessage("noext", 1, ext=".mkv"), "noext") == ".mkv" # --------------------------------------------------------------------------- # saved /rename formats # --------------------------------------------------------------------------- async def make_fmt(env, name="One Piece", **changes): fmt = await env.db.create_format(42, name, template="One Piece - E{episode}") if changes: await env.db.update_format(42, fmt["id"], **changes) return await env.db.get_format(42, fmt["id"]) @needs_ffmpeg async def test_a_format_brings_its_own_metadata_tracks_and_thumbnail(env): # the default metadata is off and has a different title: a manual format must not depend on either await env.db.update_settings(42, metadata_on=False, meta={"title": "DEFAULT"}, keep_audio=["hindi"]) await make_auto(env, thumb=b"\xff\xd8auto") # nor on the auto format fmt = await make_fmt(env, meta={"title": "One Piece", "audio": "{lang}"}, keep_audio=["english", "japanese"], keep_subtitle=["2"]) await env.db.set_format_thumb(42, fmt["id"], b"\xff\xd8format") fmt = await env.db.get_format(42, fmt["id"]) client, renamer, job = make(env) job.fmt = fmt await renamer.process(job) out = await probe(client.uploads[0]) assert [s.lang for s in out.audio] == ["eng", "jpn"] and [s.lang for s in out.subtitles] == ["eng"] assert out.tags.get("title") == "One Piece" and [s.title for s in out.audio] == ["English", "Japanese"] (_t, _f, kw), = client.sent assert kw["thumb"] == b"\xff\xd8format" assert job.state == "done" @needs_ffmpeg async def test_a_format_without_a_thumbnail_gets_a_frame_of_the_video(env): await make_auto(env, thumb=b"\xff\xd8auto") # another format's thumbnail is not borrowed fmt = await make_fmt(env) client, renamer, job = make(env) job.fmt = fmt await renamer.process(job) (_t, _f, kw), = client.sent assert kw["thumb"][:2] == b"\xff\xd8" and kw["thumb"] != b"\xff\xd8auto" # cut out of the video @needs_ffmpeg async def test_a_manual_format_that_sets_no_metadata_leaves_the_file_untouched_even_if_default_metadata_is_on(env): await env.db.update_settings(42, metadata_on=True, meta={"title": "DEFAULT"}, keep_audio=["english"]) fmt = await make_fmt(env) client, renamer, job = make(env) job.fmt = fmt await renamer.process(job) assert client.uploads[0].read_bytes() == (env.media_dir / "sample.mkv").read_bytes() # no ffmpeg pass assert job.state == "done" @needs_ffmpeg async def test_an_auto_format_without_metadata_uses_the_default_metadata_while_it_is_on(env): await env.db.update_settings(42, metadata_on=True, meta={"title": "DEFAULT TITLE"}) client, renamer, job = make(env) job.fmt = await make_auto(env) # sets no metadata of its own await renamer.process(job) assert (await probe(client.uploads[0])).tags.get("title") == "DEFAULT TITLE" await env.db.update_settings(42, metadata_on=False) # switched off: the file is left alone client, renamer, job = make(env) job.fmt = await make_auto(env) await renamer.process(job) assert client.uploads[0].read_bytes() == (env.media_dir / "sample.mkv").read_bytes() @needs_ffmpeg async def test_an_auto_format_with_metadata_of_its_own_ignores_the_default_metadata(env): await env.db.update_settings(42, metadata_on=True, meta={"title": "DEFAULT TITLE"}) client, renamer, job = make(env) job.fmt = await make_auto(env, meta={"title": "OWN TITLE"}) await renamer.process(job) assert (await probe(client.uploads[0])).tags.get("title") == "OWN TITLE" await env.db.update_settings(42, metadata_on=False) # its own metadata applies whatever the switch says client, renamer, job = make(env) job.fmt = await make_auto(env) await renamer.process(job) assert (await probe(client.uploads[0])).tags.get("title") == "OWN TITLE" @needs_ffmpeg async def test_everything_else_still_comes_from_the_users_settings(env): await env.db.update_settings(42, file_type="document", target_channel=-100777) fmt = await make_fmt(env, meta={"title": "T"}) client, renamer, job = make(env) job.fmt = fmt job.new_name = "One Piece - E5.mkv" await renamer.process(job) (target, _f, kw), = client.sent assert target == -100777 and kw["force_document"] is True # the file type and the channel are still the user's assert kw["caption"] == "One Piece - E5" # the caption is the format's: it has none, so it is the new name in bold @needs_ffmpeg async def test_a_format_sends_its_own_caption(env): from bot.utils import human_size await make_auto(env, caption="THE AUTO FORMAT'S CAPTION {filename}") # another format's caption is not borrowed fmt = await make_fmt(env, caption="{filename} · {original} · {size}") client, renamer, job = make(env) job.fmt = fmt job.new_name = "One Piece - E5.mkv" await renamer.process(job) (_t, _f, kw), = client.sent size = (env.media_dir / "sample.mkv").stat().st_size assert kw["caption"] == f"One Piece - E5 · {job.file_name} · {human_size(size)}" and kw["parse_mode"] == "html" @needs_ffmpeg async def test_a_format_without_a_caption_sends_the_new_name_in_bold(env): await make_auto(env, caption="THE AUTO FORMAT'S CAPTION {filename}") fmt = await make_fmt(env) client, renamer, job = make(env) job.fmt = fmt job.new_name = "One Piece - E5.mkv" await renamer.process(job) assert client.sent[0][2]["caption"] == "One Piece - E5" @needs_ffmpeg async def test_an_auto_rename_uses_the_caption_of_its_auto_format(env): await env.db.create_format(42, "Unused", template="X") # a manual format changes nothing for the auto renames client, renamer, job = make(env) job.fmt = await make_auto(env, caption="{filename} ({quality})") await renamer.process(job) assert client.sent[0][2]["caption"] == "[S01E05] Tokyo Revengers [1080p] (1080p)" @needs_ffmpeg async def test_an_auto_rename_without_any_format_uses_the_built_in_defaults(env): client, renamer, job = make(env) assert job.fmt is None await renamer.process(job) (_t, _f, kw), = client.sent assert kw["caption"] == "[S01E05] Tokyo Revengers [1080p]" # the new name in bold assert kw["thumb"][:2] == b"\xff\xd8" # a frame of the video assert client.uploads[0].read_bytes() == (env.media_dir / "sample.mkv").read_bytes() assert await env.db.list_formats(42) == [] # nothing was made for it @needs_ffmpeg async def test_a_format_deleted_while_queued_still_finishes_with_its_snapshot(env): fmt = await make_fmt(env, meta={"title": "One Piece"}) await env.db.set_format_thumb(42, fmt["id"], b"\xff\xd8gone") await env.db.delete_format(42, fmt["id"]) client, renamer, job = make(env) job.fmt = fmt # the snapshot taken when the file was queued await renamer.process(job) (_t, _f, kw), = client.sent assert kw["thumb"][:2] == b"\xff\xd8" and kw["thumb"] != b"\xff\xd8gone" and job.state == "done" # its thumbnail went with it out = await probe(client.uploads[0]) assert out.tags.get("title") == "One Piece" # the snapshot's metadata is still applied @needs_ffmpeg async def test_the_done_message_says_how_fast_each_way_went(env, monkeypatch): from bot import pipeline from bot.utils import human_size 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, renamer, job = make(env) await renamer.process(job) size = (env.media_dir / "sample.mkv").stat().st_size done = client.edits[-1][0] assert f"{human_size(size)} · ⬇ {human_size(size / 2)}/s · ⬆ {human_size(size / 2)}/s" in done @needs_ffmpeg async def test_a_transfer_that_took_under_a_second_shows_no_speed(env): client, renamer, job = make(env) await renamer.process(job) assert "Done" in client.edits[-1][0] and "/s" not in client.edits[-1][0]