Spaces:
Running
Running
| import asyncio | |
| import pytest | |
| from bot.errors import Cancelled, MediaError | |
| from bot.media import ( | |
| ImageTooBig, | |
| MediaInfo, | |
| Stream, | |
| apply_plan, | |
| build_plan, | |
| canonical_lang, | |
| extract_frame, | |
| format_tracks, | |
| frame_seeks, | |
| image_size, | |
| make_video_thumbnail, | |
| parse_probe, | |
| parse_track_tokens, | |
| prepare_thumb, | |
| probe, | |
| resize_exact, | |
| resolve_tokens, | |
| validate_track_tokens, | |
| ) | |
| from .conftest import ffprobe_json, needs_ffmpeg | |
| def audio(index, lang="", title=""): | |
| return Stream(index=index, kind="audio", codec="aac", lang=lang, title=title, channels=2) | |
| def sub(index, lang="", title=""): | |
| return Stream(index=index, kind="subtitle", codec="subrip", lang=lang, title=title) | |
| # --------------------------------------------------------------------------- | |
| # pure logic | |
| # --------------------------------------------------------------------------- | |
| def test_parse_track_tokens(): | |
| assert parse_track_tokens("English, Japanese") == ["english", "japanese"] | |
| assert parse_track_tokens("2 1") == ["2", "1"] | |
| assert parse_track_tokens("eng,eng jpn") == ["eng", "jpn"] | |
| assert parse_track_tokens("") == [] | |
| def test_validate_track_tokens(): | |
| assert validate_track_tokens(["2", "english", "pt-br"]) is None | |
| assert validate_track_tokens(["0"]) is not None | |
| assert validate_track_tokens(["en;rm"]) is not None | |
| def test_canonical_lang(): | |
| assert canonical_lang("eng") == "english" | |
| assert canonical_lang("EN-us") == "english" | |
| assert canonical_lang("jpn") == "japanese" | |
| assert canonical_lang("chi") == canonical_lang("zho") == "chinese" | |
| assert canonical_lang("und") is None | |
| assert canonical_lang("") is None | |
| def test_resolve_by_language_name_and_code_in_token_order(): | |
| streams = [audio(1, "jpn"), audio(2, "eng"), audio(3, "hin")] | |
| assert [s.index for s in resolve_tokens(streams, ["english", "japanese"])] == [2, 1] | |
| assert [s.index for s in resolve_tokens(streams, ["eng", "jpn"])] == [2, 1] | |
| assert [s.index for s in resolve_tokens(streams, ["hindi"])] == [3] | |
| def test_resolve_by_number_is_one_based_position_not_stream_index(): | |
| streams = [audio(5, "jpn"), audio(9, "eng")] | |
| assert [s.index for s in resolve_tokens(streams, ["2", "1"])] == [9, 5] | |
| assert resolve_tokens(streams, ["3"]) == [] | |
| def test_resolve_dedupes_and_keeps_all_matches_of_a_language(): | |
| streams = [audio(1, "eng", "Dub"), audio(2, "jpn"), audio(3, "eng", "Commentary")] | |
| assert [s.index for s in resolve_tokens(streams, ["english", "1"])] == [1, 3] | |
| assert [s.index for s in resolve_tokens(streams, ["english", "japanese"])] == [1, 3, 2] | |
| def test_resolve_untagged_track_uses_title_words(): | |
| streams = [audio(1, "und", "Japanese [Original]"), audio(2, "", "English Dub")] | |
| assert [s.index for s in resolve_tokens(streams, ["english"])] == [2] | |
| assert [s.index for s in resolve_tokens(streams, ["japanese"])] == [1] | |
| def test_resolve_tagged_language_is_not_overridden_by_title(): | |
| streams = [audio(1, "jpn", "English commentary")] | |
| assert resolve_tokens(streams, ["english"]) == [] | |
| def test_resolve_unknown_language_word_matches_raw_code(): | |
| streams = [audio(1, "xyz"), audio(2, "eng")] | |
| assert [s.index for s in resolve_tokens(streams, ["xyz"])] == [1] | |
| def make_info(): | |
| return MediaInfo( | |
| duration=2, | |
| streams=[ | |
| Stream(0, "video", codec="h264", width=320, height=240), | |
| audio(1, "jpn", "Japanese Original"), | |
| audio(2, "eng"), | |
| audio(3, "hin"), | |
| sub(4, "jpn"), | |
| sub(5, "eng"), | |
| Stream(6, "attachment", codec="ttf"), | |
| ], | |
| ) | |
| def test_build_plan_none_when_metadata_off_or_nothing_to_do(): | |
| info = make_info() | |
| assert build_plan(info, {"metadata_on": False, "keep_audio": ["1"]}, ".mkv") is None | |
| assert build_plan(info, {"metadata_on": True, "meta": {}}, ".mkv") is None | |
| assert build_plan(info, {"metadata_on": True, "keep_audio": ["1"]}, ".avi") is None | |
| def test_build_plan_maps_in_requested_order_and_drops_others(): | |
| plan = build_plan( | |
| make_info(), | |
| {"metadata_on": True, "keep_audio": ["english", "japanese"], "keep_subtitle": ["2"], "meta": {}}, | |
| ".mkv", | |
| ) | |
| assert plan is not None | |
| maps = [plan.args[i + 1] for i, a in enumerate(plan.args) if a == "-map"] | |
| # video, audio eng(2) then jpn(1), subtitle #2 -> stream 5, attachment kept; hindi + jpn sub dropped | |
| assert maps == ["0:0", "0:2", "0:1", "0:5", "0:6"] | |
| assert plan.args[plan.args.index("-disposition:a:0") + 1] == "+default" | |
| assert plan.args[plan.args.index("-disposition:a:1") + 1] == "-default" | |
| assert plan.args[plan.args.index("-disposition:s:0") + 1] == "+default" | |
| def test_build_plan_audio_safeguard(): | |
| plan = build_plan(make_info(), {"metadata_on": True, "keep_audio": ["klingon"], "meta": {}}, ".mkv") | |
| assert plan is not None | |
| assert [s.index for s in plan.kept_audio] == [1, 2, 3] | |
| assert plan.warnings | |
| def test_build_plan_subtitles_dropped_when_nothing_matches(): | |
| plan = build_plan(make_info(), {"metadata_on": True, "keep_subtitle": ["french"], "meta": {}}, ".mkv") | |
| assert plan.kept_subtitles == [] | |
| maps = [plan.args[i + 1] for i, a in enumerate(plan.args) if a == "-map"] | |
| assert "0:4" not in maps and "0:5" not in maps | |
| def test_build_plan_metadata_fields_and_lang_placeholder(): | |
| plan = build_plan( | |
| make_info(), | |
| { | |
| "metadata_on": True, | |
| "keep_audio": ["japanese", "english"], | |
| "meta": {"title": "Tokyo Revengers", "author": "@chan", "artist": "@chan", | |
| "audio": "{lang} Audio", "subtitle": "English", "video": "@chan"}, | |
| }, | |
| ".mkv", | |
| ) | |
| a = plan.args | |
| assert ["-metadata", "title=Tokyo Revengers"] == a[a.index("title=Tokyo Revengers") - 1: a.index("title=Tokyo Revengers") + 1] | |
| assert "author=@chan" in a and "artist=@chan" in a | |
| assert "title=Japanese Audio" in a and "title=English Audio" in a | |
| assert a.index("title=Japanese Audio") < a.index("title=English Audio") | |
| assert "-metadata:s:v:0" in a and "-metadata:s:s:0" in a | |
| def test_build_plan_mp3_tags_only(): | |
| plan = build_plan(MediaInfo(), {"metadata_on": True, "meta": {"title": "T", "author": "A"}}, ".mp3") | |
| assert plan is not None | |
| assert "title=T" in plan.args and "artist=A" in plan.args | |
| assert build_plan(MediaInfo(), {"metadata_on": True, "meta": {}}, ".mp3") is None | |
| def test_format_tracks(): | |
| text = format_tracks(make_info()) | |
| assert "1. Japanese (jpn)" in text and "2. English (eng)" in text and "3. Hindi (hin)" in text | |
| assert "Subtitles" in text and "“Japanese Original”" in text | |
| def test_parse_probe_duration_fallbacks(): | |
| info = parse_probe({"format": {"duration": "12.5"}, "streams": []}) | |
| assert info.duration == 12.5 | |
| info = parse_probe({"format": {}, "streams": [{"index": 0, "codec_type": "video", "tags": {"DURATION": "00:00:03.500000000"}}]}) | |
| assert info.duration == 3.5 | |
| # --------------------------------------------------------------------------- | |
| # real ffprobe / ffmpeg | |
| # --------------------------------------------------------------------------- | |
| async def test_probe_real_mkv(media_dir): | |
| info = await probe(media_dir / "sample.mkv") | |
| assert [s.kind for s in info.streams] == ["video", "audio", "audio", "audio", "subtitle", "subtitle", "attachment"] | |
| assert [s.lang for s in info.audio] == ["jpn", "eng", "hin"] | |
| assert [s.lang for s in info.subtitles] == ["jpn", "eng"] | |
| assert info.width == 320 and info.height == 240 | |
| assert 1.5 < info.duration < 2.6 | |
| async def test_probe_missing_file_raises(tmp_path): | |
| with pytest.raises(MediaError): | |
| await probe(tmp_path / "nope.mkv") | |
| async def test_apply_plan_mkv_reorders_drops_and_tags(media_dir, tmp_path): | |
| src = media_dir / "sample.mkv" | |
| info = await probe(src) | |
| settings = { | |
| "metadata_on": True, | |
| "keep_audio": ["english", "japanese"], | |
| "keep_subtitle": ["english"], | |
| "meta": {"title": "Tokyo Revengers", "author": "@chan", "artist": "@chan", | |
| "audio": "{lang} Audio", "subtitle": "{lang}", "video": "@chan video"}, | |
| } | |
| plan = build_plan(info, settings, ".mkv") | |
| dst = tmp_path / "out.mkv" | |
| await apply_plan(src, dst, plan) | |
| data = ffprobe_json(dst, "-show_chapters") | |
| streams = data["streams"] | |
| kinds = [s["codec_type"] for s in streams] | |
| assert kinds == ["video", "audio", "audio", "subtitle", "attachment"] | |
| audios = [s for s in streams if s["codec_type"] == "audio"] | |
| assert [a["tags"]["language"] for a in audios] == ["eng", "jpn"] # reordered, hindi dropped | |
| assert [a["tags"]["title"] for a in audios] == ["English Audio", "Japanese Audio"] | |
| assert audios[0]["disposition"]["default"] == 1 and audios[1]["disposition"]["default"] == 0 | |
| sub_stream = next(s for s in streams if s["codec_type"] == "subtitle") | |
| assert sub_stream["tags"]["language"] == "eng" and sub_stream["tags"]["title"] == "English" | |
| assert sub_stream["disposition"]["default"] == 1 | |
| video = streams[0] | |
| assert video["tags"]["title"] == "@chan video" | |
| fmt_tags = {k.lower(): v for k, v in data["format"]["tags"].items()} # Matroska stores AUTHOR/ARTIST upper-case | |
| assert fmt_tags["title"] == "Tokyo Revengers" | |
| assert fmt_tags["author"] == "@chan" and fmt_tags["artist"] == "@chan" | |
| assert len(data["chapters"]) == 2 # chapters survive | |
| attach = next(s for s in streams if s["codec_type"] == "attachment") | |
| assert attach["tags"]["filename"] == "font.ttf" # fonts survive | |
| async def test_apply_plan_numbers_follow_given_order(media_dir, tmp_path): | |
| src = media_dir / "sample.mkv" | |
| info = await probe(src) | |
| plan = build_plan(info, {"metadata_on": True, "keep_audio": ["3", "1"], "keep_subtitle": ["2", "1"], "meta": {}}, ".mkv") | |
| dst = tmp_path / "out.mkv" | |
| await apply_plan(src, dst, plan) | |
| out = await probe(dst) | |
| assert [s.lang for s in out.audio] == ["hin", "jpn"] | |
| assert [s.lang for s in out.subtitles] == ["eng", "jpn"] | |
| async def test_apply_plan_audio_safeguard_keeps_everything(media_dir, tmp_path): | |
| src = media_dir / "sample.mkv" | |
| info = await probe(src) | |
| plan = build_plan(info, {"metadata_on": True, "keep_audio": ["french"], "meta": {}}, ".mkv") | |
| dst = tmp_path / "out.mkv" | |
| await apply_plan(src, dst, plan) | |
| assert len((await probe(dst)).audio) == 3 | |
| async def test_apply_plan_mp4(media_dir, tmp_path): | |
| src = media_dir / "sample.mp4" | |
| info = await probe(src) | |
| plan = build_plan(info, {"metadata_on": True, "keep_audio": ["english", "japanese"], "keep_subtitle": ["1"], | |
| "meta": {"title": "MP4 Title"}}, ".mp4") | |
| dst = tmp_path / "out.mp4" | |
| await apply_plan(src, dst, plan) | |
| out = await probe(dst) | |
| assert [s.lang for s in out.audio] == ["eng", "jpn"] | |
| assert [s.kind for s in out.streams].count("subtitle") == 1 | |
| assert out.tags.get("title") == "MP4 Title" | |
| async def test_apply_plan_mp3_title(media_dir, tmp_path): | |
| src = media_dir / "sample.mp3" | |
| plan = build_plan(await probe(src), {"metadata_on": True, "meta": {"title": "New Title", "artist": "Me"}}, ".mp3") | |
| dst = tmp_path / "out.mp3" | |
| await apply_plan(src, dst, plan) | |
| out = await probe(dst) | |
| assert out.tags["title"] == "New Title" and out.tags["artist"] == "Me" | |
| async def test_apply_plan_failure_raises_media_error(tmp_path): | |
| bad = tmp_path / "bad.mkv" | |
| bad.write_bytes(b"not a video") | |
| from bot.media import Plan | |
| with pytest.raises(MediaError): | |
| await apply_plan(bad, tmp_path / "out.mkv", Plan(["-map", "0:0", "-c", "copy"])) | |
| async def test_cancel_kills_ffmpeg(media_dir, tmp_path): | |
| cancel = asyncio.Event() | |
| cancel.set() | |
| info = await probe(media_dir / "sample.mkv") | |
| plan = build_plan(info, {"metadata_on": True, "keep_audio": ["1"], "meta": {}}, ".mkv") | |
| with pytest.raises(Cancelled): | |
| await apply_plan(media_dir / "sample.mkv", tmp_path / "out.mkv", plan, cancel) | |
| async def test_thumbnail_and_prepare_thumb(media_dir, tmp_path): | |
| out = tmp_path / "t.jpg" | |
| assert await make_video_thumbnail(media_dir / "sample.mkv", out, 2.0) | |
| assert out.read_bytes()[:2] == b"\xff\xd8" | |
| from PIL import Image | |
| import io | |
| big = io.BytesIO() | |
| Image.effect_noise((1600, 1200), 80).convert("RGB").save(big, "PNG") | |
| jpg = prepare_thumb(big.getvalue()) | |
| img = Image.open(io.BytesIO(jpg)) | |
| assert max(img.size) <= 320 and len(jpg) <= 200 * 1024 and img.format == "JPEG" | |
| assert not await make_video_thumbnail(tmp_path / "missing.mkv", tmp_path / "x.jpg", 1.0) | |
| # --------------------------------------------------------------------------- | |
| # /ethumb, /size and /resize helpers | |
| # --------------------------------------------------------------------------- | |
| def test_frame_seeks_start_at_about_a_tenth_and_fall_back_to_earlier_positions(): | |
| assert frame_seeks(1400) == [60.0, 20.0, 5.0, 0.0] | |
| assert frame_seeks(100) == [10.0, 5.0, 0.0] | |
| assert frame_seeks(2) == [0.2, 0.0] | |
| assert frame_seeks(0) == [1.0, 0.0] # duration unknown (a cut-off file) | |
| assert frame_seeks(-5) == [1.0, 0.0] | |
| def _png(size, color=(10, 120, 200), mode="RGB", exif_orientation=None): | |
| import io | |
| from PIL import Image | |
| img = Image.new(mode, size, color) | |
| kw = {} | |
| if exif_orientation: | |
| exif = img.getexif() | |
| exif[274] = exif_orientation | |
| kw["exif"] = exif | |
| buf = io.BytesIO() | |
| img.save(buf, "PNG", **kw) | |
| return buf.getvalue() | |
| def _opened(data): | |
| import io | |
| from PIL import Image | |
| return Image.open(io.BytesIO(data)) | |
| def test_image_size_is_the_width_and_the_height(): | |
| assert image_size(_png((1920, 1080))) == (1920, 1080) | |
| assert image_size(_png((7, 300))) == (7, 300) | |
| def test_image_size_counts_a_picture_stored_sideways_as_turned_upright(): | |
| assert image_size(_png((200, 100), exif_orientation=6)) == (100, 200) | |
| assert image_size(_png((200, 100), exif_orientation=8)) == (100, 200) | |
| assert image_size(_png((200, 100), exif_orientation=1)) == (200, 100) | |
| assert image_size(_png((200, 100), exif_orientation=3)) == (200, 100) # upside down: the same width and height | |
| def test_image_size_rejects_what_is_not_an_image(): | |
| with pytest.raises(Exception): | |
| image_size(b"definitely not an image") | |
| def test_resize_exact_gives_exactly_the_size_asked_for_whatever_the_shape(source, wanted): | |
| img = _opened(resize_exact(_png(source), *wanted)) | |
| assert img.format == "PNG" and img.size == wanted | |
| def test_resize_exact_stretches_the_whole_picture(): | |
| from PIL import Image | |
| two_halves = Image.new("RGB", (200, 100), (0, 0, 255)) | |
| two_halves.paste((255, 0, 0), (0, 0, 100, 100)) | |
| import io | |
| buf = io.BytesIO() | |
| two_halves.save(buf, "PNG") | |
| out = _opened(resize_exact(buf.getvalue(), 40, 400)).convert("RGB") | |
| assert out.getpixel((2, 5)) == out.getpixel((2, 395)) == (255, 0, 0) | |
| assert out.getpixel((37, 5)) == out.getpixel((37, 395)) == (0, 0, 255) | |
| def test_resize_exact_keeps_transparency_and_does_not_turn_it_white_or_black(): | |
| img = _opened(resize_exact(_png((200, 200), (0, 0, 0, 0), mode="RGBA"), 100, 50)) | |
| assert img.mode == "RGBA" and img.getpixel((50, 25))[3] == 0 | |
| def test_resize_exact_turns_a_picture_stored_sideways_upright_first(): | |
| from PIL import Image | |
| import io | |
| img = Image.new("RGB", (200, 100), (0, 0, 255)) | |
| img.paste((255, 0, 0), (0, 0, 100, 100)) # the left half is red | |
| exif = img.getexif() | |
| exif[274] = 6 | |
| buf = io.BytesIO() | |
| img.save(buf, "PNG", exif=exif) | |
| out = _opened(resize_exact(buf.getvalue(), 100, 200)).convert("RGB") | |
| assert out.getpixel((50, 10)) == (255, 0, 0) and out.getpixel((50, 190)) == (0, 0, 255) # red on top: it was turned | |
| def test_resize_exact_takes_every_kind_of_picture(mode, saved_as): | |
| import io | |
| from PIL import Image | |
| buf = io.BytesIO() | |
| Image.new(mode, (30, 20)).save(buf, saved_as) | |
| img = _opened(resize_exact(buf.getvalue(), 15, 10)) | |
| assert img.format == "PNG" and img.size == (15, 10) | |
| def test_resize_exact_keeps_a_palette_picture_smooth_not_blocky(): | |
| import io | |
| from PIL import Image | |
| ramp = Image.new("L", (256, 1)) | |
| ramp.putdata(list(range(256))) | |
| buf = io.BytesIO() | |
| ramp.convert("P").save(buf, "PNG") | |
| out = _opened(resize_exact(buf.getvalue(), 64, 1)).convert("L") | |
| values = [out.getpixel((x, 0)) for x in range(64)] | |
| assert len(set(values)) > 50 and values == sorted(values) # a gradient, not a few flat bands | |
| def test_resize_exact_refuses_a_picture_with_too_many_pixels(monkeypatch): | |
| monkeypatch.setattr("bot.media.RESIZE_MAX_PIXELS", 400) | |
| with pytest.raises(ImageTooBig): | |
| resize_exact(_png((21, 20)), 5, 5) | |
| assert _opened(resize_exact(_png((20, 20)), 5, 5)).size == (5, 5) # right at the limit is fine | |
| assert issubclass(ImageTooBig, ValueError) | |
| def test_resize_exact_rejects_what_is_not_an_image(): | |
| with pytest.raises(Exception): | |
| resize_exact(b"definitely not an image", 100, 100) | |
| async def test_extract_frame_writes_a_jpeg_and_tries_the_next_position_after_a_miss(media_dir, tmp_path): | |
| out = tmp_path / "f.jpg" | |
| assert await extract_frame(media_dir / "sample.mkv", out, [1.0]) | |
| assert out.read_bytes()[:2] == b"\xff\xd8" | |
| from PIL import Image | |
| assert Image.open(out).size == (320, 240) # the sample is smaller than the 1280 px limit | |
| out.unlink() | |
| assert await extract_frame(media_dir / "sample.mkv", out, [500.0, 0.0]) # nothing at 500 s: the next position works | |
| assert out.stat().st_size > 0 | |
| async def test_extract_frame_scales_down_wide_pictures(media_dir, tmp_path): | |
| from PIL import Image | |
| out = tmp_path / "small.jpg" | |
| assert await extract_frame(media_dir / "sample.mkv", out, [0.0], width=160) | |
| assert Image.open(out).size == (160, 120) | |
| async def test_extract_frame_gives_up_on_junk_without_raising(tmp_path): | |
| junk = tmp_path / "junk.mkv" | |
| junk.write_bytes(b"\x00" * 4096) | |
| assert not await extract_frame(junk, tmp_path / "j.jpg", [1.0, 0.0]) | |
| assert not await extract_frame(tmp_path / "missing.mkv", tmp_path / "m.jpg", [0.0]) | |