Add Embedded Subtitles provider

3 files changed, 115 insertions, 0 deletions

diff --git a/tests/subliminal_patch/data/file_1.mkv b/tests/subliminal_patch/data/file_1.mkv
new file mode 100644
index 000000000..61112e0c2
--- /dev/null
+++ b/tests/subliminal_patch/data/file_1.mkv
diff --git a/tests/subliminal_patch/data/file_2.mkv b/tests/subliminal_patch/data/file_2.mkv
new file mode 100644
index 000000000..eefabffe9
--- /dev/null
+++ b/tests/subliminal_patch/data/file_2.mkv
diff --git a/tests/subliminal_patch/test_embeddedsubtitles.py b/tests/subliminal_patch/test_embeddedsubtitles.py
new file mode 100644
index 000000000..59ffaf0f2
--- /dev/null
+++ b/tests/subliminal_patch/test_embeddedsubtitles.py
@@ -0,0 +1,115 @@
+# -*- coding: utf-8 -*-
+import os
+
+import fese
+import pytest
+from subliminal_patch.core import Episode, Movie
+from subliminal_patch.providers.embeddedsubtitles import EmbeddedSubtitlesProvider
+from subzero.language import Language
+
+_DATA = os.path.join(os.path.abspath(os.path.dirname(__file__)), "data")
+
+
+fese.Language = Language
+
+
[email protected]
+def video_single_language():
+    # Has only ASS streams in english
+    return Episode(
+        os.path.join(_DATA, "file_1.mkv"),
+        "Serial Experiments Lain",
+        1,
+        1,
+        source="Web",
+    )
+
+
[email protected]
+def video_multiple_languages():
+    # Has SubRip streams in multiple languages
+    return Movie(
+        os.path.join(_DATA, "file_2.mkv"),
+        "I'm No Longer Here",
+        year=2019,
+        source="Web",
+    )
+
+
[email protected]
+def video_inexistent(tmpdir):
+    return Movie(
+        os.path.join(tmpdir, "inexistent_video.mkv"),
+        "Dummy",
+        year=2021,
+        source="Web",
+    )
+
+
+def test_inexistent_video(video_inexistent):
+    with EmbeddedSubtitlesProvider() as provider:
+        subtitles = provider.list_subtitles(video_inexistent, {})
+        assert len(subtitles) == 0
+
+
+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_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"
+        }
+        with pytest.raises(fese.InvalidFile):
+            provider.download_subtitle(subtitle)