Spaces:
Running
Running
| 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)) | |
| 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 | |
| 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"] == "<b>[S01E05] Tokyo Revengers [1080p]</b>" 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 | |
| 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] | |
| 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" | |
| 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"]} | |
| 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"]} | |
| 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"]) | |
| 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" | |
| 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] | |
| 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" | |
| 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] | |
| 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" | |
| 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() | |
| 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() | |
| 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'<pre><code class="language-python">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</code></pre>") | |
| assert not (Path(env.cfg.temp_dir) / "job1").exists() | |
| 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 == [] | |
| 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] | |
| 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"]) | |
| 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" | |
| 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 | |
| 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" | |
| 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() | |
| 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" | |
| 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"] == "<b>One Piece - E5</b>" # the caption is the format's: it has none, so it is the new name in bold | |
| async def test_a_format_sends_its_own_caption(env): | |
| from bot.utils import human_size | |
| await make_auto(env, caption="<i>THE AUTO FORMAT'S CAPTION {filename}</i>") # another format's caption is not borrowed | |
| fmt = await make_fmt(env, caption="<b>{filename}</b> · {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"<b>One Piece - E5</b> · {job.file_name} · {human_size(size)}" and kw["parse_mode"] == "html" | |
| async def test_a_format_without_a_caption_sends_the_new_name_in_bold(env): | |
| await make_auto(env, caption="<i>THE AUTO FORMAT'S CAPTION {filename}</i>") | |
| 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"] == "<b>One Piece - E5</b>" | |
| 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="<i>{filename}</i> ({quality})") | |
| await renamer.process(job) | |
| assert client.sent[0][2]["caption"] == "<i>[S01E05] Tokyo Revengers [1080p]</i> (1080p)" | |
| 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"] == "<b>[S01E05] Tokyo Revengers [1080p]</b>" # 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 | |
| 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 | |
| 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 | |
| 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] | |