bazarr/tests/subliminal_patch/test_embeddedsubtitles.py

# -*- coding: utf-8 -*-
import os

import fese
import pytest
import fese
from fese import FFprobeSubtitleStream
from subliminal_patch.core import Episode, Movie
from subliminal_patch.exceptions import MustGetBlacklisted
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


@pytest.fixture
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",
    )


@pytest.fixture
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",
    )


@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"},
            }
        ),
    }


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",
            return_value=[fake_streams["en_hi"]],
        )
        fake = fese.FFprobeVideoContainer.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",
            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_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",
            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