# -*- coding: utf-8 -*- import os import tempfile import fese from fese import FFprobeSubtitleStream import pytest import subliminal_patch from subliminal_patch.core import Episode from subliminal_patch.core import Movie from subliminal_patch.exceptions import MustGetBlacklisted from subliminal_patch.providers.embeddedsubtitles import _MemoizedFFprobeVideoContainer from subliminal_patch.providers.embeddedsubtitles import EmbeddedSubtitlesProvider from subliminal_patch.providers.embeddedsubtitles import ( _discard_possible_incomplete_subtitles, ) from subzero.language import Language fese.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 { "include_ass": True, "include_srt": True, "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_init(config): with EmbeddedSubtitlesProvider(**config) as provider: assert provider is not None 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"}, } ), } 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( # "fese.FFprobeVideoContainer.get_subtitles", "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 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].hearing_impaired == True assert subs[1].hearing_impaired == 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(include_ass=False) 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(include_srt=False) as provider: subs = provider.list_subtitles( video_multiple_languages, {Language.fromalpha2("en")} ) assert not subs 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.content is not None 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.content is not None def test_download_invalid_subtitle(video_single_language): with EmbeddedSubtitlesProvider() as provider: subtitle = provider.list_subtitles( video_single_language, {Language.fromalpha2("en")} )[0] provider._cached_paths[subtitle.container.path] = { subtitle.stream.index: "dummy.srt" } try: provider.download_subtitle(subtitle) except MustGetBlacklisted as error: assert error.id == subtitle.id assert error.media_type == subtitle.media_type 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