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 # --------------------------------------------------------------------------- @needs_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 @needs_ffmpeg async def test_probe_missing_file_raises(tmp_path): with pytest.raises(MediaError): await probe(tmp_path / "nope.mkv") @needs_ffmpeg 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 @needs_ffmpeg 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"] @needs_ffmpeg 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 @needs_ffmpeg 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" @needs_ffmpeg 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" @needs_ffmpeg 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"])) @needs_ffmpeg 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) @needs_ffmpeg 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") @pytest.mark.parametrize("source,wanted", [ ((1920, 1080), (720, 1080)), ((100, 50), (400, 400)), ((1000, 2000), (3, 900)), ((640, 480), (1, 1)), ((64, 64), (64, 64)), ]) 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 @pytest.mark.parametrize("mode,saved_as", [("P", "PNG"), ("1", "PNG"), ("LA", "PNG"), ("L", "PNG"), ("CMYK", "JPEG"), ("RGBA", "PNG")]) 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) @needs_ffmpeg 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 @needs_ffmpeg 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) @needs_ffmpeg 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])