[DEV] Remaining throttle unit tests (#1075)

This commit is contained in:
Jesse Bannon 2024-09-30 22:04:32 -07:00 committed by GitHub
parent 1caade46bd
commit 01e1f46b86
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 85 additions and 111 deletions

View file

@ -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

View file

@ -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 (

View file

@ -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,

View file

@ -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

View file

@ -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