# -*- coding: utf-8 -*- import os from fese import FFprobeSubtitleStream from fese import FFprobeVideoContainer from fese import tags from fese.exceptions import LanguageNotFound import pytest from subliminal_patch.core import Episode from subliminal_patch.core import Movie from subliminal_patch.providers.embeddedsubtitles import ( _discard_possible_incomplete_subtitles, _clean_ass_subtitles, ) from subliminal_patch.providers.embeddedsubtitles import _get_pretty_release_name from subliminal_patch.providers.embeddedsubtitles import _MemoizedFFprobeVideoContainer from subliminal_patch.providers.embeddedsubtitles import EmbeddedSubtitlesProvider from subzero.language import Language tags.Language = Language @pytest.fixture def video_single_language(data): # Has only ASS streams in english return Episode( os.path.join(data, "file_1.mkv"), "Serial Experiments Lain", 1, 1, source="Web", ) @pytest.fixture def video_multiple_languages(data): # Has SubRip streams in multiple languages return Movie( os.path.join(data, "file_2.mkv"), "I'm No Longer Here", year=2019, source="Web", ) @pytest.fixture def config(tmpdir): return { "included_codecs": None, "cache_dir": tmpdir, "ffprobe_path": None, "ffmpeg_path": None, "hi_fallback": False, } @pytest.fixture def video_inexistent(tmpdir): return Movie( os.path.join(tmpdir, "inexistent_video.mkv"), "Dummy", year=2021, source="Web", ) def test_language_is_subzero_type(): assert tags.Language == Language def test_init(config): with EmbeddedSubtitlesProvider(**config) as provider: assert provider is not None def test_init_empty_included_codecs(): with EmbeddedSubtitlesProvider(included_codecs=[]) as provider: assert provider._included_codecs == {"ass", "subrip", "webvtt", "mov_text"} def test_init_custom_included_codecs(): with EmbeddedSubtitlesProvider(included_codecs=["ass"]) as provider: assert provider._included_codecs == {"ass"} def test_inexistent_video(video_inexistent): with EmbeddedSubtitlesProvider() as provider: subtitles = provider.list_subtitles(video_inexistent, {}) assert len(subtitles) == 0 @pytest.fixture def fake_streams(): return { "en_hi": FFprobeSubtitleStream( { "index": 3, "codec_name": "subrip", "disposition": {"default": 1, "hearing_impaired": 1}, "tags": {"language": "eng", "title": "English"}, } ), "en": FFprobeSubtitleStream( { "index": 3, "codec_name": "subrip", "tags": {"language": "eng", "title": "English"}, } ), "es_hi": FFprobeSubtitleStream( { "index": 3, "codec_name": "subrip", "disposition": {"default": 1, "hearing_impaired": 1}, "tags": {"language": "spa", "title": "Spanish"}, } ), "es": FFprobeSubtitleStream( { "index": 3, "codec_name": "subrip", "tags": {"language": "spa", "title": "Spanish"}, } ), } @pytest.mark.parametrize("tags_", [{}, {"language": "und", "title": "Unknown"}]) def test_list_subtitles_unknown_as_english(mocker, tags_, video_single_language): with EmbeddedSubtitlesProvider(unknown_as_english=True) as provider: fake = FFprobeSubtitleStream( {"index": 3, "codec_name": "subrip", "tags": tags_} ) mocker.patch( "subliminal_patch.providers.embeddedsubtitles._MemoizedFFprobeVideoContainer.get_subtitles", return_value=[fake], ) result = provider.list_subtitles( video_single_language, {Language.fromalpha2("en")} ) assert len(result) == 1 def test_list_subtitles_unknown_as_english_w_real_english_subtitles( video_single_language, mocker ): with EmbeddedSubtitlesProvider(unknown_as_english=True) as provider: fakes = [ FFprobeSubtitleStream( {"index": 3, "codec_name": "subrip", "tags": {"language": "und"}} ), FFprobeSubtitleStream( {"index": 2, "codec_name": "subrip", "tags": {"language": "eng"}} ), ] mocker.patch( "subliminal_patch.providers.embeddedsubtitles._MemoizedFFprobeVideoContainer.get_subtitles", return_value=fakes, ) result = provider.list_subtitles( video_single_language, {Language.fromalpha2("en")} ) assert len(result) == 1 @pytest.mark.parametrize("tags_", [{}, {"language": "und", "title": "Unknown"}]) def test_list_subtitles_unknown_as_english_disabled(tags_): with EmbeddedSubtitlesProvider(unknown_as_english=False): with pytest.raises(LanguageNotFound): assert FFprobeSubtitleStream( {"index": 3, "codec_name": "subrip", "tags": tags_} ) def test_list_subtitles_hi_fallback_one_stream( video_single_language, fake_streams, mocker ): with EmbeddedSubtitlesProvider(hi_fallback=True) as provider: language = Language.fromalpha2("en") mocker.patch( "subliminal_patch.providers.embeddedsubtitles._MemoizedFFprobeVideoContainer.get_subtitles", return_value=[fake_streams["en_hi"]], ) fake = _MemoizedFFprobeVideoContainer.get_subtitles("")[0] assert fake.disposition.hearing_impaired == True subs = provider.list_subtitles(video_single_language, {language}) assert subs[0].language == Language("eng", hi=False) assert subs[0].hearing_impaired == False def test_list_subtitles_hi_fallback_multiple_streams( video_single_language, fake_streams, mocker ): with EmbeddedSubtitlesProvider(hi_fallback=True) as provider: language = Language.fromalpha2("en") mocker.patch( # "fese.FFprobeVideoContainer.get_subtitles", "subliminal_patch.providers.embeddedsubtitles._MemoizedFFprobeVideoContainer.get_subtitles", return_value=[fake_streams["en_hi"], fake_streams["en"]], ) subs = provider.list_subtitles(video_single_language, {language}) assert len(subs) == 2 assert subs[0].language == Language("eng", hi=True) assert subs[1].language == Language("eng", hi=False) def test_list_subtitles_hi_fallback_multiple_language_streams( video_single_language, fake_streams, mocker ): with EmbeddedSubtitlesProvider(hi_fallback=True) as provider: languages = {Language.fromalpha2("en"), Language.fromalpha2("es")} mocker.patch( # "fese.FFprobeVideoContainer.get_subtitles", "subliminal_patch.providers.embeddedsubtitles._MemoizedFFprobeVideoContainer.get_subtitles", return_value=[ fake_streams["en_hi"], fake_streams["es"], fake_streams["es_hi"], ], ) subs = provider.list_subtitles(video_single_language, languages) assert len(subs) == 3 assert subs[0].hearing_impaired == False # English subittle assert subs[1].hearing_impaired == False # Spanish subtitle assert subs[2].hearing_impaired == True # Spanish HI subtitle def test_list_subtitles_hi_fallback_multiple_hi_streams( video_single_language, fake_streams, mocker ): with EmbeddedSubtitlesProvider(hi_fallback=True) as provider: language = Language.fromalpha2("en") mocker.patch( # "fese.FFprobeVideoContainer.get_subtitles", "subliminal_patch.providers.embeddedsubtitles._MemoizedFFprobeVideoContainer.get_subtitles", return_value=[fake_streams["en_hi"], fake_streams["en_hi"]], ) subs = provider.list_subtitles(video_single_language, {language}) assert len(subs) == 2 assert subs[0].hearing_impaired == False assert subs[1].hearing_impaired == False def test_list_subtitles_only_forced(video_single_language): with EmbeddedSubtitlesProvider() as provider: language = Language.fromalpha2("en") language = Language.rebuild(language, forced=True) subs = provider.list_subtitles(video_single_language, {language}) assert len(subs) == 0 def test_list_subtitles_also_forced(video_single_language): with EmbeddedSubtitlesProvider() as provider: language_1 = Language.fromalpha2("en") language_2 = Language.rebuild(language_1, forced=True) subs = provider.list_subtitles(video_single_language, {language_1, language_2}) assert any(language_1 == sub.language for sub in subs) assert any(not sub.language.forced for sub in subs) def test_list_subtitles_single_language(video_single_language): with EmbeddedSubtitlesProvider() as provider: subs = provider.list_subtitles( video_single_language, {Language.fromalpha2("en")} ) for sub in subs: assert sub.language == Language.fromalpha2("en") def test_list_subtitles_multiple_languages(video_multiple_languages): with EmbeddedSubtitlesProvider() as provider: languages = {Language.fromalpha2(code) for code in ("en", "it", "fr", "es")} | { Language("por", "BR") } subs = provider.list_subtitles(video_multiple_languages, languages) for expected in languages: assert any(sub.language == expected for sub in subs) def test_list_subtitles_wo_ass(video_single_language): with EmbeddedSubtitlesProvider(included_codecs=("srt",)) as provider: subs = provider.list_subtitles( video_single_language, {Language.fromalpha2("en")} ) assert not subs def test_list_subtitles_wo_srt(video_multiple_languages): with EmbeddedSubtitlesProvider(included_codecs=("ass",)) as provider: subs = provider.list_subtitles( video_multiple_languages, {Language.fromalpha2("en")} ) assert not subs def test_get_pretty_release_name(): stream = FFprobeSubtitleStream( { "index": 1, "codec_name": "subrip", "tags": {"language": "eng", "title": "forced"}, } ) container = FFprobeVideoContainer("foo.mkv") assert _get_pretty_release_name(stream, container) == "foo.en.forced.srt" def test_clean_ass_subtitles(data, tmp_path): path = os.path.join(data, "subs.ass") with open(path, "r") as f: og_lines_len = len(f.readlines()) output_path = os.path.join(tmp_path, "subs.ass") _clean_ass_subtitles(path, output_path) with open(output_path, "r") as f: assert og_lines_len > len(f.readlines()) def test_download_subtitle_multiple(video_multiple_languages): with EmbeddedSubtitlesProvider() as provider: languages = {Language.fromalpha2(code) for code in ("en", "it", "fr")} | { Language("por", "BR") } subs = provider.list_subtitles(video_multiple_languages, languages) for sub in subs: provider.download_subtitle(sub) assert sub.is_valid() def test_download_subtitle_single(video_single_language): with EmbeddedSubtitlesProvider() as provider: subtitle = provider.list_subtitles( video_single_language, {Language.fromalpha2("en")} )[0] provider.download_subtitle(subtitle) assert subtitle.is_valid() def test_memoized(video_single_language, mocker): with EmbeddedSubtitlesProvider() as provider: provider.list_subtitles(video_single_language, {Language.fromalpha2("en")}) with EmbeddedSubtitlesProvider() as provider: mocker.patch("fese.FFprobeVideoContainer.get_subtitles") assert ( provider.list_subtitles(video_single_language, {Language.fromalpha2("en")})[ 0 ] is not None ) @pytest.mark.parametrize( "number_of_frames,expected_len", [((34, 811), 1), ((0, 0), 2), ((811, 34), 1), ((900, 1000), 2), ((0, 900), 1)], ) def test_discard_possible_incomplete_subtitles(number_of_frames, expected_len): subtitle_1 = FFprobeSubtitleStream( { "index": 1, "codec_name": "subrip", "codec_long_name": "SubRip subtitle", "disposition": {}, "tags": {"language": "eng", "NUMBER_OF_FRAMES": number_of_frames[0]}, } ) subtitle_2 = FFprobeSubtitleStream( { "index": 2, "codec_name": "subrip", "codec_long_name": "SubRip subtitle", "disposition": {}, "tags": {"language": "eng", "NUMBER_OF_FRAMES": number_of_frames[1]}, } ) new_list = _discard_possible_incomplete_subtitles([subtitle_1, subtitle_2]) assert len(new_list) == expected_len