SuperSubtitles Provider: fix filetype detection

3 files changed, 120 insertions, 9 deletions

diff --git a/libs/subliminal_patch/providers/supersubtitles.py b/libs/subliminal_patch/providers/supersubtitles.py
index 01aff5a29..a37852344 100644
--- a/libs/subliminal_patch/providers/supersubtitles.py
+++ b/libs/subliminal_patch/providers/supersubtitles.py
@@ -18,9 +18,10 @@ from subliminal.providers import ParserBeautifulSoup
 from bs4.element import Tag, NavigableString
 from subliminal.score import get_equivalent_release_groups
 from subliminal_patch.subtitle import Subtitle, guess_matches
+from subliminal_patch.exceptions import APIThrottled
 from subliminal.utils import sanitize, sanitize_release_group
 from subliminal.video import Episode, Movie
-from zipfile import ZipFile
+from zipfile import ZipFile, is_zipfile
 from rarfile import RarFile, is_rarfile
 from subliminal_patch.utils import sanitize, fix_inconsistent_naming
 from guessit import guessit
@@ -197,6 +198,7 @@ class SuperSubtitlesProvider(Provider, ProviderSubtitleArchiveMixin):
         """
 
         """
+        # TODO: add memoization to this method logic
 
         url = self.server_url + "index.php?tipus=adatlap&azon=a_" + str(sub_id)
         # url = https://www.feliratok.info/index.php?tipus=adatlap&azon=a_1518600916
@@ -212,6 +214,7 @@ class SuperSubtitlesProvider(Provider, ProviderSubtitleArchiveMixin):
                 # src="img/adatlap/imdb.png"/></a>
                 imdb_id = re.search(r'(?<=www\.imdb\.com/title/).*(?=/")', str(value))
                 imdb_id = imdb_id.group() if imdb_id else ''
+                logger.debug("IMDB ID found: %s", imdb_id)
                 return imdb_id
 
         return None
@@ -546,15 +549,15 @@ class SuperSubtitlesProvider(Provider, ProviderSubtitleArchiveMixin):
         r = self.session.get(subtitle.page_link, timeout=10)
         r.raise_for_status()
 
-        if ".rar" in subtitle.page_link:
-            logger.debug('Archive identified as rar')
-            archive_stream = io.BytesIO(r.content)
+        archive_stream = io.BytesIO(r.content)
+        archive = None
+
+        if is_rarfile(archive_stream):
             archive = RarFile(archive_stream)
-            subtitle.content = self.get_subtitle_from_archive(subtitle, archive)
-        elif ".zip" in subtitle.page_link:
-            logger.debug('Archive identified as zip')
-            archive_stream = io.BytesIO(r.content)
+        elif is_zipfile(archive_stream):
             archive = ZipFile(archive_stream)
-            subtitle.content = self.get_subtitle_from_archive(subtitle, archive)
         else:
             subtitle.content = fix_line_ending(r.content)
+
+        if archive is not None:
+            subtitle.content = self.get_subtitle_from_archive(subtitle, archive)
diff --git a/tests/subliminal_patch/conftest.py b/tests/subliminal_patch/conftest.py
index e9c539a62..167acfad1 100644
--- a/tests/subliminal_patch/conftest.py
+++ b/tests/subliminal_patch/conftest.py
@@ -27,6 +27,7 @@ def movies():
             resolution="1080p",
             source="Web",
             # other="Rip",
+            imdb_id="tt1160419",
             alternative_titles=["Dune: Part One"],
             audio_codec="Dolby Digital",
             video_codec="H.264",
diff --git a/tests/subliminal_patch/test_supersubtitles.py b/tests/subliminal_patch/test_supersubtitles.py
new file mode 100644
index 000000000..6111cabc0
--- /dev/null
+++ b/tests/subliminal_patch/test_supersubtitles.py
@@ -0,0 +1,107 @@
+# -*- coding: utf-8 -*-
+
+import pytest
+from subliminal_patch.providers.supersubtitles import SuperSubtitlesProvider
+from subliminal_patch.providers.supersubtitles import SuperSubtitlesSubtitle
+from subliminal_patch.core import Episode
+from subzero.language import Language
+
+
[email protected]
+def episode():
+    episode = {
+        "name": "/tv/All of Us Are Dead/Season 1/All of Us Are Dead - S01E11 - Episode 11 WEBDL-1080p.mp4",
+        "source": "Web",
+        "release_group": None,
+        "resolution": "1080p",
+        "video_codec": None,
+        "audio_codec": None,
+        "imdb_id": None,
+        "subtitle_languages": set(),
+        "streaming_service": None,
+        "edition": None,
+        "series": "All of Us Are Dead",
+        "season": 1,
+        "episode": 11,
+        "title": "Episode 11",
+        "year": None,
+        "original_series": True,
+        "tvdb_id": None,
+        "series_tvdb_id": None,
+        "series_imdb_id": None,
+        "alternative_series": [],
+    }
+    return Episode(**episode)
+
+
+def test_list_episode_subtitles(episode):
+    language = Language.fromalpha2("en")
+
+    with SuperSubtitlesProvider() as provider:
+        assert provider.list_subtitles(episode, {language})
+
+
+def test_download_episode_subtitle(episode):
+    subtitle = SuperSubtitlesSubtitle(
+        Language.fromalpha2("en"),
+        "https://www.feliratok.info/index.php?action=letolt&felirat=1643361676",
+        1643361676,
+        "All of us are dead",
+        1,
+        11,
+        "",
+        [
+            "NF.WEB-DL.1080p-TEPES",
+            "NF.WEBRip.1080p-TEPES",
+            "WEBRip-ION10",
+            "WEBRip-ION265",
+            "WEBRip.1080p-RARBG",
+        ],
+        "",
+        "",
+        "",
+        asked_for_episode=True,
+    )
+    assert subtitle.get_matches(episode)
+
+    with SuperSubtitlesProvider() as provider:
+        provider.download_subtitle(subtitle)
+        assert subtitle.is_valid()
+
+
+def test_list_and_download_movie_subtitles(movies):
+    movie = movies["dune"]
+    language = Language.fromalpha2("en")
+
+    with SuperSubtitlesProvider() as provider:
+        assert provider.list_subtitles(movie, {language})
+
+
+def test_download_movie_subtitle(movies):
+    movie = movies["dune"]
+
+    subtitle = SuperSubtitlesSubtitle(
+        Language.fromalpha2("en"),
+        "https://www.feliratok.info/index.php?action=letolt&felirat=1634579718",
+        1634579718,
+        "Dune",
+        0,
+        0,
+        "",
+        [
+            "NF.WEB-DL.1080p-TEPES",
+            "NF.WEBRip.1080p-TEPES",
+            "WEBRip-ION10",
+            "WEBRip-ION265",
+            "WEBRip.1080p-RARBG",
+        ],
+        "",
+        "",
+        "",
+        asked_for_episode=None,
+    )
+    assert subtitle.get_matches(movie)
+
+    with SuperSubtitlesProvider() as provider:
+        provider.download_subtitle(subtitle)
+        assert subtitle.is_valid()