From 949fd26b24a1350017c72e5ecd57944381095b53 Mon Sep 17 00:00:00 2001 From: vitiko98 Date: Thu, 30 Jun 2022 18:07:45 -0400 Subject: [PATCH] Start subliminal_patch's core refactor --- libs/subliminal_patch/core.py | 22 +--- tests/subliminal_patch/conftest.py | 4 +- tests/subliminal_patch/data/dummy.srt | 40 +++++++ tests/subliminal_patch/test_core.py | 152 ++++++++++++++++++++++++++ 4 files changed, 196 insertions(+), 22 deletions(-) create mode 100644 tests/subliminal_patch/data/dummy.srt create mode 100644 tests/subliminal_patch/test_core.py diff --git a/libs/subliminal_patch/core.py b/libs/subliminal_patch/core.py index 806ec0bba..fe9c38bb6 100644 --- a/libs/subliminal_patch/core.py +++ b/libs/subliminal_patch/core.py @@ -1,6 +1,5 @@ # coding=utf-8 from __future__ import absolute_import -import codecs import json import re import os @@ -13,7 +12,6 @@ import operator import unicodedata import itertools -from six.moves.http_client import ResponseNotReady import rarfile import requests @@ -27,13 +25,11 @@ from concurrent.futures import as_completed from .extensions import provider_registry from .exceptions import MustGetBlacklisted -from subliminal.exceptions import ServiceUnavailable, DownloadLimitExceeded from subliminal.score import compute_score as default_compute_score from subliminal.utils import hash_napiprojekt, hash_opensubtitles, hash_shooter, hash_thesubdb from subliminal.video import VIDEO_EXTENSIONS, Video, Episode, Movie from subliminal.core import guessit, ProviderPool, io, is_windows_special_path, \ ThreadPoolExecutor, check_video -from subliminal_patch.exceptions import TooManyRequests, APIThrottled from subzero.language import Language, ENDSWITH_LANGUAGECODE_RE, FULL_LANGUAGE_LIST try: @@ -142,8 +138,7 @@ class _Blacklist(list): class SZProviderPool(ProviderPool): - def __init__(self, providers=None, provider_configs=None, blacklist=None, ban_list=None, throttle_callback=None, - pre_download_hook=None, post_download_hook=None, language_hook=None): + def __init__(self, providers=None, provider_configs=None, blacklist=None, ban_list=None, throttle_callback=None): #: Name of providers to use self.providers = set(providers or []) @@ -160,10 +155,6 @@ class SZProviderPool(ProviderPool): self.throttle_callback = throttle_callback - self.pre_download_hook = pre_download_hook - self.post_download_hook = post_download_hook - self.language_hook = language_hook - self._born = time.time() if not self.throttle_callback: @@ -266,10 +257,7 @@ class SZProviderPool(ProviderPool): :rtype: list of :class:`~subliminal.subtitle.Subtitle` or None """ - if self.language_hook: - languages_search_base = self.language_hook(provider) - else: - languages_search_base = languages + languages_search_base = languages # check video validity if not provider_registry[provider].check(video): @@ -377,13 +365,7 @@ class SZProviderPool(ProviderPool): while True: tries += 1 try: - if self.pre_download_hook: - self.pre_download_hook(subtitle) - self[subtitle.provider_name].download_subtitle(subtitle) - if self.post_download_hook: - self.post_download_hook(subtitle) - break except (requests.ConnectionError, requests.exceptions.ProxyError, diff --git a/tests/subliminal_patch/conftest.py b/tests/subliminal_patch/conftest.py index d60a807aa..c8d499f57 100644 --- a/tests/subliminal_patch/conftest.py +++ b/tests/subliminal_patch/conftest.py @@ -17,7 +17,7 @@ def vcr_cassette_dir(request): return os.path.join("tests/subliminal_patch/cassettes", request.module.__name__) -@pytest.fixture +@pytest.fixture(scope="session") def movies(): return { "dune": Movie( @@ -85,7 +85,7 @@ def movies(): } -@pytest.fixture +@pytest.fixture(scope="session") def episodes(): return { "got_s03e10": Episode( diff --git a/tests/subliminal_patch/data/dummy.srt b/tests/subliminal_patch/data/dummy.srt new file mode 100644 index 000000000..bd2d9c936 --- /dev/null +++ b/tests/subliminal_patch/data/dummy.srt @@ -0,0 +1,40 @@ +1 +00:00:00,000 --> 00:00:05,000 +Foo. + +2 +00:00:05,000 --> 00:00:07,000 +Bar. + +3 +00:00:07,000 --> 00:00:09,000 +Foo. + +4 +00:00:09,000 --> 00:00:11,000 +Bar. + +5 +00:00:11,000 --> 00:00:13,000 +Foo. + +6 +00:00:13,000 --> 00:00:15,000 +Bar. + +7 +00:00:15,000 --> 00:00:17,000 +Foo. + +8 +00:00:17,000 --> 00:00:19,000 +Bar. + +9 +00:00:19,000 --> 00:00:21,000 +Foo. + +10 +00:00:21,000 --> 00:00:23,000 +Bar + diff --git a/tests/subliminal_patch/test_core.py b/tests/subliminal_patch/test_core.py new file mode 100644 index 000000000..6adeed256 --- /dev/null +++ b/tests/subliminal_patch/test_core.py @@ -0,0 +1,152 @@ +import os + +import pytest +from subliminal_patch import core +from subliminal_patch.score import compute_score +from subliminal_patch import core_persistent as corep +from subliminal_patch import Provider +from subliminal_patch import Subtitle +from subliminal_patch.core import SZProviderPool as Pool +from subzero.language import Language + + +class FakeProviderSubtitle(Subtitle): + provider_name = "fake" + + def __init__(self, language, id, matches=None): + super().__init__(language, page_link=id) + self._id = id + self._matches = set(matches or []) + + self.release_info = id + + def get_matches(self, video): + return self._matches + + def id(self): + return self.id + + +_ENGLISH = Language.fromietf("en") + + +class FakeProvider(Provider): + languages = {_ENGLISH} + video_types = (core.Movie, core.Episode) + + def __init__(self, fake_subtitles=None): + self._fake_subtitles = fake_subtitles + + def initialize(self): + pass + + def terminate(self): + pass + + def list_subtitles(self, video, languages): + return self._fake_subtitles + + def download_subtitle(self, subtitle): + filename = os.path.join( + os.path.abspath(os.path.dirname(__file__)), "data", "dummy.srt" + ) + with open(filename, "rb") as f: + subtitle.content = f.read() + + +core.provider_registry.register("fake", FakeProvider) + + +def test_pool_init_default(): + with Pool() as pool: + assert pool + + +def test_pool_init_w_providers(): + with Pool( + providers={"opensubtitles"}, + provider_configs={"opensubtitles": {"username": "foo", "password": "bar"}}, + ) as pool: + assert pool.providers is not None + assert isinstance(pool.provider_configs, core._ProviderConfigs) + + +def test_pool_init_w_blacklist(): + with Pool(blacklist=[("foo", "bar")]) as pool: + assert isinstance(pool.blacklist, core._Blacklist) + + +def test_pool_init_w_ban_list(): + with Pool(ban_list={"must_contain": ["foo"], "must_not_contain": ["bar"]}) as pool: + assert isinstance(pool.ban_list, core._Banlist) + + +def test_pool_init_w_throttle_callback(): + with Pool(throttle_callback=None) as pool: + pool.throttle_callback("foo", "bar") + + +@pytest.fixture(scope="module") +def fake_subtitle(): + yield FakeProviderSubtitle(_ENGLISH, "foo") + + +@pytest.fixture(scope="module") +def fake_pool(fake_subtitle): + with Pool( + providers={"fake"}, + provider_configs={"fake": {"fake_subtitles": [fake_subtitle]}}, + ) as pool: + yield pool + + +def test_pool_list_subtitles_provider(fake_pool, fake_subtitle, movies): + result = fake_pool.list_subtitles_provider("fake", movies["dune"], {_ENGLISH}) + assert fake_subtitle in result + + +def test_pool_list_subtitles(fake_pool, fake_subtitle, movies): + result = fake_pool.list_subtitles(movies["dune"], {_ENGLISH}) + assert fake_subtitle in result + + +def test_pool_download_subtitle(fake_pool, fake_subtitle): + downloaded = fake_pool.download_subtitle(fake_subtitle) + assert downloaded is True + + +def test_pool_download_best_subtitles(fake_pool, fake_subtitle, movies): + result = fake_pool.download_best_subtitles( + [fake_subtitle], movies["dune"], [_ENGLISH], compute_score=compute_score + ) + assert result == [fake_subtitle] + + +@pytest.fixture(scope="module") +def empty_pool(): + with Pool() as pool: + yield pool + + +def test_pool_core_persistent_list_all_subtitles(movies, empty_pool): + assert corep.list_all_subtitles([movies["dune"]], {_ENGLISH}, empty_pool) == { + movies["dune"]: [] + } + + +def test_pool_core_persistent_list_supported_languages(empty_pool): + assert corep.list_supported_languages(empty_pool) == [] + + +def test_pool_core_persistent_list_supported_video_types(empty_pool): + assert corep.list_supported_video_types(empty_pool) == [] + + +def test_pool_core_persistent_download_subtitles(empty_pool): + corep.download_subtitles([], empty_pool) + + +def test_pool_core_persistent_download_best_subtitles(movies, empty_pool): + assert corep.download_best_subtitles([movies["dune"]], {_ENGLISH}, empty_pool) == { + movies["dune"]: [] + }