From 89e387030dd861384298636f18ad668595d3e3fa Mon Sep 17 00:00:00 2001 From: Patrick Robertson Date: Tue, 18 Mar 2025 09:59:59 +0000 Subject: [PATCH] Tests for suitable URLs for tikwm --- .../generic_extractor/generic_extractor.py | 2 ++ .../modules/generic_extractor/tiktok.py | 2 +- tests/conftest.py | 2 +- .../extractors/test_tiktok_tikwm_extractor.py | 28 ++++++++++++++++++- 4 files changed, 31 insertions(+), 3 deletions(-) diff --git a/src/auto_archiver/modules/generic_extractor/generic_extractor.py b/src/auto_archiver/modules/generic_extractor/generic_extractor.py index e7d6be0..80556bf 100644 --- a/src/auto_archiver/modules/generic_extractor/generic_extractor.py +++ b/src/auto_archiver/modules/generic_extractor/generic_extractor.py @@ -75,6 +75,8 @@ class GenericExtractor(Extractor): dropin: GenericDropin = self.dropin_for_name(info_extractor.ie_key()) if dropin and dropin.suitable(url, info_extractor): yield info_extractor + elif info_extractor.suitable(url): + yield info_extractor def suitable(self, url: str) -> bool: """ diff --git a/src/auto_archiver/modules/generic_extractor/tiktok.py b/src/auto_archiver/modules/generic_extractor/tiktok.py index e545ba9..e44714e 100644 --- a/src/auto_archiver/modules/generic_extractor/tiktok.py +++ b/src/auto_archiver/modules/generic_extractor/tiktok.py @@ -16,7 +16,7 @@ class Tiktok(GenericDropin): TIKWM_ENDPOINT = "https://www.tikwm.com/api/?url={url}" - def suitable(self, url, info_extractor): + def suitable(self, url, info_extractor) -> bool: """This dropin (which uses Tikvm) is suitable for *all* Tiktok type URLs - videos, lives, VMs, and users. Return the 'suitable' method from the TikTokIE class.""" return any(extractor().suitable(url) for extractor in (TikTokIE, TikTokLiveIE, TikTokVMIE, TikTokUserIE)) diff --git a/tests/conftest.py b/tests/conftest.py index 379bfc2..6e87e26 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -118,7 +118,7 @@ def pytest_runtest_setup(item): pytest.xfail(f"previous test failed ({test_name})") -@pytest.fixture() +@pytest.fixture def unpickle(): """ Returns a helper function that unpickles a file diff --git a/tests/extractors/test_tiktok_tikwm_extractor.py b/tests/extractors/test_tiktok_tikwm_extractor.py index d04d7e4..81f29a5 100644 --- a/tests/extractors/test_tiktok_tikwm_extractor.py +++ b/tests/extractors/test_tiktok_tikwm_extractor.py @@ -4,6 +4,8 @@ import pytest import yt_dlp from auto_archiver.modules.generic_extractor.generic_extractor import GenericExtractor +from auto_archiver.modules.generic_extractor.tiktok import Tiktok, TikTokIE + from .test_extractor_base import TestExtractorBase @@ -17,11 +19,16 @@ def skip_ytdlp_own_methods(mocker): ) -@pytest.fixture() +@pytest.fixture def mock_get(mocker): return mocker.patch("auto_archiver.modules.generic_extractor.tiktok.requests.get") +@pytest.fixture +def tiktok_dropin() -> Tiktok: + return Tiktok() + + class TestTiktokTikwmExtractor(TestExtractorBase): """ Test suite for TestTiktokTikwmExtractor. @@ -34,6 +41,25 @@ class TestTiktokTikwmExtractor(TestExtractorBase): VALID_EXAMPLE_URL = "https://www.tiktok.com/@example/video/1234" + @pytest.mark.parametrize( + "url, is_suitable", + [ + ("https://bellingcat.com", False), + ("https://youtube.com", False), + ("https://tiktok.co/", False), + ("https://tiktok.com/", False), + ("https://www.tiktok.com/", False), + ("https://api.cool.tiktok.com/", False), + (VALID_EXAMPLE_URL, True), + ("https://www.tiktok.com/@bbcnews/video/7478038212070411542", True), + ("https://www.tiktok.com/@ggs68taiwan.official/video/7441821351142362375", True), + ("https://www.tiktok.com/t/ZP8YQ8e5j/", True), + ("https://vt.tiktok.com/ZSMTJeqRP/", True), + ], + ) + def test_is_suitable(self, url, is_suitable, tiktok_dropin): + assert tiktok_dropin.suitable(url, TikTokIE()) == is_suitable + def test_invalid_json_responses(self, mock_get, make_item, caplog): mock_get.return_value.status_code = 200 mock_get.return_value.json.side_effect = ValueError