diff --git a/src/ytdl_sub/config/plugin/plugin.py b/src/ytdl_sub/config/plugin/plugin.py index e3c6d1d8..13feb41d 100644 --- a/src/ytdl_sub/config/plugin/plugin.py +++ b/src/ytdl_sub/config/plugin/plugin.py @@ -69,6 +69,13 @@ class Plugin(BasePlugin[OptionsValidatorT], Generic[OptionsValidatorT], ABC): ytdl options to enable/disable when downloading entries for this specific plugin """ + def initialize_subscription(self) -> bool: + """ + Before any downloading begins, perform initialization before the subscription runs. + Returns true if this subscription should run, false otherwise. + """ + return True + def modify_entry_metadata(self, entry: Entry) -> Optional[Entry]: """ After entry metadata has been gathered, perform preprocessing on the metadata diff --git a/src/ytdl_sub/plugins/throttle_protection.py b/src/ytdl_sub/plugins/throttle_protection.py index 9fb102db..5bd45f10 100644 --- a/src/ytdl_sub/plugins/throttle_protection.py +++ b/src/ytdl_sub/plugins/throttle_protection.py @@ -1,8 +1,6 @@ import random import time -from typing import List from typing import Optional -from typing import Tuple from ytdl_sub.config.overrides import Overrides from ytdl_sub.config.plugin.plugin import Plugin @@ -167,29 +165,17 @@ class ThrottleProtectionPlugin(Plugin[ThrottleProtectionOptions]): self.plugin_options.max_downloads_per_subscription.randomized_int() ) - def ytdl_options_match_filters(self) -> Tuple[List[str], List[str]]: - """ - Returns - ------- - If subscription_download_probability, match-filters that will perform no downloads - if it's rolled to not download. - """ - perform_download: Tuple[List[str], List[str]] = [], [] - do_not_perform_download: Tuple[List[str], List[str]] = [], [ - "title = __YTDL_SUB_THROTTLE_PROTECTION_ON_SUBSCRIPTION_DOWNLOAD__" - ] - + def initialize_subscription(self) -> bool: if self.plugin_options.subscription_download_probability: proba = self.plugin_options.subscription_download_probability.value # assume proba is set to 1.0, random.random() will always be < 1, can never reach this if random.random() > proba: logger.info( - "Subscription download probability of %f missed, skipping this subscription", + "Subscription download probability of %0.2f missed", proba, ) - return do_not_perform_download - - return perform_download + return False + return True def modify_entry_metadata(self, entry: Entry) -> Optional[Entry]: if ( diff --git a/src/ytdl_sub/subscriptions/subscription_download.py b/src/ytdl_sub/subscriptions/subscription_download.py index 790ffd26..5d1b4cb3 100644 --- a/src/ytdl_sub/subscriptions/subscription_download.py +++ b/src/ytdl_sub/subscriptions/subscription_download.py @@ -348,6 +348,10 @@ class SubscriptionDownload(BaseSubscription, ABC): self.download_archive.reinitialize(dry_run=dry_run) plugins = self._initialize_plugins() + if not all(plugin.initialize_subscription() for plugin in plugins): + # Any plugin that skips gracefully should have logs that explain why + logging.info("Skipping %s", self.name) + return FileHandlerTransactionLog() subscription_ytdl_options = SubscriptionYTDLOptions( preset=self._preset_options, diff --git a/tests/e2e/plugins/test_throttle_protection.py b/tests/e2e/plugins/test_throttle_protection.py deleted file mode 100644 index 92ef4c2a..00000000 --- a/tests/e2e/plugins/test_throttle_protection.py +++ /dev/null @@ -1,35 +0,0 @@ -import pytest - -from ytdl_sub.subscriptions.subscription import Subscription - - -@pytest.fixture -def preset_dict_subscription_download_proba_0(output_directory): - return { - "preset": "Jellyfin Music Videos", - "download": "https://youtube.com/watch?v=HKTNxEqsN3Q", - "format": "worst[ext=mp4]", - "overrides": { - "music_video_artist": "JMC", - "music_video_directory": output_directory, - }, - "throttle_protection": {"subscription_download_probability": 0.0}, - } - - -class TestThrottleProtection: - - def test_subscription_probability( - self, - default_config, - preset_dict_subscription_download_proba_0, - output_directory, - ): - single_video_subscription = Subscription.from_dict( - config=default_config, - preset_name="music_video_single_video_test", - preset_dict=preset_dict_subscription_download_proba_0, - ) - - transaction_log = single_video_subscription.download(dry_run=True) - assert transaction_log.is_empty diff --git a/tests/unit/plugins/test_throttle_protection.py b/tests/unit/plugins/test_throttle_protection.py index 47e02a2b..424887ad 100644 --- a/tests/unit/plugins/test_throttle_protection.py +++ b/tests/unit/plugins/test_throttle_protection.py @@ -1,3 +1,5 @@ +from typing import Dict + import pytest from conftest import assert_logs @@ -5,38 +7,44 @@ from ytdl_sub.plugins.throttle_protection import logger as throttle_protection_l from ytdl_sub.subscriptions.subscription import Subscription +@pytest.fixture +def throttle_subscription_dict(output_directory) -> Dict: + return { + "preset": [ + "Kodi Music Videos", + ], + "overrides": { + "url": "https://your.name.here", + "music_video_directory": output_directory, + "bool_false_variable": "{ %bool(False) }", + "empty_string_variable": "", + }, + "throttle_protection": { + "sleep_per_download_s": { + "min": 0.01, + "max": 0.01, + }, + "sleep_per_subscription_s": { + "min": 0.02, + "max": 0.02, + }, + }, + } + + class TestThrottleProtectionPlugin: def test_sleeps_log( self, config, subscription_name, + throttle_subscription_dict, output_directory, mock_download_collection_entries, ): - preset_dict = { - "preset": [ - "Kodi Music Videos", - ], - "overrides": { - "url": "https://your.name.here", - "music_video_directory": output_directory, - }, - "throttle_protection": { - "sleep_per_download_s": { - "min": 0.01, - "max": 0.01, - }, - "sleep_per_subscription_s": { - "min": 0.02, - "max": 0.02, - }, - }, - } - subscription = Subscription.from_dict( config=config, preset_name=subscription_name, - preset_dict=preset_dict, + preset_dict=throttle_subscription_dict, ) with ( @@ -70,45 +78,24 @@ class TestThrottleProtectionPlugin: [ "", False, - "{tp_bool_string}", - "{tp_empty_string}", + "{bool_false_variable}", + "{empty_string_variable}", ], ) def test_disabled( self, config, subscription_name, + throttle_subscription_dict, output_directory, mock_download_collection_entries, disable_value, ): - preset_dict = { - "preset": [ - "Kodi Music Videos", - ], - "overrides": { - "url": "https://your.name.here", - "music_video_directory": output_directory, - "tp_bool_string": "{ %bool(False) }", - "tp_empty_string": "", - }, - "throttle_protection": { - "enable": disable_value, - "sleep_per_download_s": { - "min": 0.01, - "max": 0.01, - }, - "sleep_per_subscription_s": { - "min": 0.02, - "max": 0.02, - }, - }, - } - + throttle_subscription_dict["throttle_protection"]["enable"] = disable_value subscription = Subscription.from_dict( config=config, preset_name=subscription_name, - preset_dict=preset_dict, + preset_dict=throttle_subscription_dict, ) with ( @@ -128,24 +115,17 @@ class TestThrottleProtectionPlugin: self, config, subscription_name, + throttle_subscription_dict, output_directory, mock_download_collection_entries, ): - preset_dict = { - "preset": [ - "Kodi Music Videos", - ], - "overrides": { - "url": "https://your.name.here", - "music_video_directory": output_directory, - }, - "throttle_protection": {"max_downloads_per_subscription": {"max": 0}}, + throttle_subscription_dict["throttle_protection"] = { + "max_downloads_per_subscription": {"max": 0} } - subscription = Subscription.from_dict( config=config, preset_name=subscription_name, - preset_dict=preset_dict, + preset_dict=throttle_subscription_dict, ) with ( @@ -162,3 +142,35 @@ class TestThrottleProtectionPlugin: transaction_log = subscription.download(dry_run=True) assert transaction_log.is_empty + + def test_subscription_proba( + self, + config, + subscription_name, + throttle_subscription_dict, + output_directory, + mock_download_collection_entries, + ): + throttle_subscription_dict["throttle_protection"] = { + "subscription_download_probability": 0.0 + } + subscription = Subscription.from_dict( + config=config, + preset_name=subscription_name, + preset_dict=throttle_subscription_dict, + ) + + with ( + mock_download_collection_entries( + is_youtube_channel=False, num_urls=1, is_extracted_audio=False, is_dry_run=True + ), + assert_logs( + logger=throttle_protection_logger, + expected_message="Subscription download probability of", + log_level="info", + expected_occurrences=1, + ), + ): + transaction_log = subscription.download(dry_run=True) + + assert transaction_log.is_empty