This commit is contained in:
Jesse Bannon 2023-11-04 22:46:41 -07:00
parent a1e2f134e8
commit a932dad87b
10 changed files with 190 additions and 31 deletions

View file

@ -17,6 +17,7 @@ from ytdl_sub.plugins.output_directory_nfo_tags import OutputDirectoryNfoTagsPlu
from ytdl_sub.plugins.regex import RegexPlugin from ytdl_sub.plugins.regex import RegexPlugin
from ytdl_sub.plugins.split_by_chapters import SplitByChaptersPlugin from ytdl_sub.plugins.split_by_chapters import SplitByChaptersPlugin
from ytdl_sub.plugins.subtitles import SubtitlesPlugin from ytdl_sub.plugins.subtitles import SubtitlesPlugin
from ytdl_sub.plugins.throttle_protection import ThrottleProtectionPlugin
from ytdl_sub.plugins.video_tags import VideoTagsPlugin from ytdl_sub.plugins.video_tags import VideoTagsPlugin
@ -41,6 +42,7 @@ class PluginMapping:
"subtitles": SubtitlesPlugin, "subtitles": SubtitlesPlugin,
"chapters": ChaptersPlugin, "chapters": ChaptersPlugin,
"split_by_chapters": SplitByChaptersPlugin, "split_by_chapters": SplitByChaptersPlugin,
"throttle_protection": ThrottleProtectionPlugin,
} }
@classmethod @classmethod

View file

@ -12,7 +12,7 @@ from mergedeep import mergedeep
from ytdl_sub.config.config_validator import ConfigValidator from ytdl_sub.config.config_validator import ConfigValidator
from ytdl_sub.config.plugin import Plugin from ytdl_sub.config.plugin import Plugin
from ytdl_sub.config.preset_class_mappings import PluginMapping from ytdl_sub.config.plugin_mapping import PluginMapping
from ytdl_sub.config.preset_options import OptionsValidator from ytdl_sub.config.preset_options import OptionsValidator
from ytdl_sub.config.preset_options import OutputOptions from ytdl_sub.config.preset_options import OutputOptions
from ytdl_sub.config.preset_options import Overrides from ytdl_sub.config.preset_options import Overrides

View file

@ -34,12 +34,14 @@ class RandomizedRangeValidator(StrictDictValidator):
key="min", validator=FloatValidator, default=0.0 key="min", validator=FloatValidator, default=0.0
).value ).value
if self._min < 0:
raise self._validation_exception("min must be greater than zero")
if self._max < self._min: if self._max < self._min:
raise self._validation_exception( raise self._validation_exception(
f"max ({self._max}) must be greater than or equal to min ({self._min})" f"max ({self._max}) must be greater than or equal to min ({self._min})"
) )
@property
def randomized_float(self) -> float: def randomized_float(self) -> float:
""" """
Returns Returns
@ -48,14 +50,13 @@ class RandomizedRangeValidator(StrictDictValidator):
""" """
return random.uniform(self._min, self._max) return random.uniform(self._min, self._max)
@property
def randomized_int(self) -> int: def randomized_int(self) -> int:
""" """
Returns Returns
------- -------
A random integer within the range after casting the min + max to ints A random float within the range, then cast to an integer (floored)
""" """
return random.randrange(int(self._min), int(self._max)) return int(self.randomized_float())
class ThrottleProtectionOptions(OptionsDictValidator): class ThrottleProtectionOptions(OptionsDictValidator):
@ -124,7 +125,7 @@ class ThrottleProtectionOptions(OptionsDictValidator):
""" """
Range of downloads to perform per subscription. Range of downloads to perform per subscription.
""" """
return self._sleep_per_subscription_s return self._max_downloads_per_subscription
@property @property
def subscription_download_probability(self) -> Optional[ProbabilityValidator]: def subscription_download_probability(self) -> Optional[ProbabilityValidator]:
@ -152,7 +153,7 @@ class ThrottleProtectionPlugin(Plugin[ThrottleProtectionOptions]):
# If subscriptions have a max download limit, set it here for the first subscription # If subscriptions have a max download limit, set it here for the first subscription
if self.plugin_options.max_downloads_per_subscription: if self.plugin_options.max_downloads_per_subscription:
self._subscription_max_downloads = ( self._subscription_max_downloads = (
self.plugin_options.max_downloads_per_subscription.randomized_int self.plugin_options.max_downloads_per_subscription.randomized_int()
) )
def ytdl_options_match_filters(self) -> Tuple[List[str], List[str]]: def ytdl_options_match_filters(self) -> Tuple[List[str], List[str]]:
@ -168,8 +169,13 @@ class ThrottleProtectionPlugin(Plugin[ThrottleProtectionOptions]):
] ]
if self.plugin_options.subscription_download_probability: if self.plugin_options.subscription_download_probability:
# assume proba is set to 1.0, random.random() will always be < 1, so do nothing proba = self.plugin_options.subscription_download_probability.value
if random.random() < 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",
proba,
)
return do_not_perform_download return do_not_perform_download
return perform_download return perform_download
@ -181,7 +187,8 @@ class ThrottleProtectionPlugin(Plugin[ThrottleProtectionOptions]):
): ):
if self._subscription_download_counter == self._subscription_max_downloads: if self._subscription_download_counter == self._subscription_max_downloads:
logger.info( logger.info(
"reached subscription max downloads of %d", self._subscription_max_downloads "Reached subscription max downloads of %d for throttle protection",
self._subscription_max_downloads,
) )
self._subscription_download_counter += 1 # increment to only print once self._subscription_download_counter += 1 # increment to only print once
@ -190,17 +197,20 @@ class ThrottleProtectionPlugin(Plugin[ThrottleProtectionOptions]):
return entry return entry
def post_process_entry(self, entry: Entry) -> Optional[FileMetadata]: def post_process_entry(self, entry: Entry) -> Optional[FileMetadata]:
if self._subscription_download_counter == 0: if (
logger.info( self._subscription_max_downloads is not None
"setting subscription max downloads to %d", self._subscription_max_downloads and self._subscription_download_counter == 0
):
logger.debug(
"Setting subscription max downloads to %d", self._subscription_max_downloads
) )
# Increment the counter # Increment the counter
self._subscription_download_counter += 1 self._subscription_download_counter += 1
if self.plugin_options.sleep_per_download_s: if self.plugin_options.sleep_per_download_s:
sleep_time = self.plugin_options.sleep_per_download_s.randomized_float sleep_time = self.plugin_options.sleep_per_download_s.randomized_float()
logger.info("sleeping between downloads for %0.2f seconds", sleep_time) logger.debug("Sleeping between downloads for %0.2f seconds", sleep_time)
time.sleep(sleep_time) time.sleep(sleep_time)
return None return None
@ -216,6 +226,6 @@ class ThrottleProtectionPlugin(Plugin[ThrottleProtectionOptions]):
) )
if self.plugin_options.sleep_per_subscription_s: if self.plugin_options.sleep_per_subscription_s:
sleep_time = self.plugin_options.sleep_per_subscription_s.randomized_float sleep_time = self.plugin_options.sleep_per_subscription_s.randomized_float()
logger.info("sleeping between subscriptions for %0.2f seconds", sleep_time) logger.debug("Sleeping between subscriptions for %0.2f seconds", sleep_time)
time.sleep(sleep_time) time.sleep(sleep_time)

View file

@ -9,6 +9,7 @@ from typing import Any
from typing import Callable from typing import Callable
from typing import Dict from typing import Dict
from typing import List from typing import List
from typing import Optional
from unittest.mock import patch from unittest.mock import patch
import pytest import pytest
@ -78,11 +79,17 @@ def reformat_directory() -> Path:
@contextlib.contextmanager @contextlib.contextmanager
def assert_logs(logger: logging.Logger, expected_message: str, log_level: str = "debug"): def assert_logs(
logger: logging.Logger,
expected_message: str,
log_level: str = "debug",
expected_occurrences: Optional[int] = None,
):
""" """
Patches any function, but calls the original function. Patches any function, but calls the original function.
Intended to see if the particular function is called. Intended to see if the particular function is called.
""" """
occurrences = 0
debug_logger = Logger.get() debug_logger = Logger.get()
def _wrapped_debug(*args, **kwargs): def _wrapped_debug(*args, **kwargs):
@ -92,10 +99,14 @@ def assert_logs(logger: logging.Logger, expected_message: str, log_level: str =
yield yield
for call_args in patched_debug.call_args_list: for call_args in patched_debug.call_args_list:
if expected_message in call_args.args[0]: occurrences += int(expected_message in call_args.args[0])
return
assert False, f"{expected_message} was not found in a logger.debug call" if expected_occurrences:
assert (
occurrences == expected_occurrences
), f"{expected_message} was expected {expected_occurrences} times, got {occurrences}"
else:
assert occurrences > 0, f"{expected_message} was not found in a logger.debug call"
def preset_dict_to_dl_args(preset_dict: Dict) -> str: def preset_dict_to_dl_args(preset_dict: Dict) -> str:

View file

@ -0,0 +1,72 @@
import pytest
from conftest import assert_logs
from ytdl_sub.plugins.throttle_protection import logger as throttle_protection_logger
from ytdl_sub.subscriptions.subscription import Subscription
@pytest.fixture
def preset_dict_max_downloads_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": {"max_downloads_per_subscription": {"max": 0}},
}
@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_max_downloads(
self,
default_config,
preset_dict_max_downloads_0,
output_directory,
):
single_video_subscription = Subscription.from_dict(
config=default_config,
preset_name="music_video_single_video_test",
preset_dict=preset_dict_max_downloads_0,
)
with assert_logs(
logger=throttle_protection_logger,
expected_message="Reached subscription max downloads of %d",
log_level="info",
expected_occurrences=1,
):
transaction_log = single_video_subscription.download(dry_run=True)
assert transaction_log.is_empty
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

@ -21,7 +21,6 @@ def single_video_preset_dict_old_video_tags_format(output_directory):
"download": "https://youtube.com/watch?v=HKTNxEqsN3Q", "download": "https://youtube.com/watch?v=HKTNxEqsN3Q",
# override the output directory with our fixture-generated dir # override the output directory with our fixture-generated dir
"output_options": { "output_options": {
"output_directory": output_directory,
"maintain_download_archive": False, "maintain_download_archive": False,
}, },
"embed_thumbnail": True, # embed thumb into the video "embed_thumbnail": True, # embed thumb into the video
@ -32,7 +31,10 @@ def single_video_preset_dict_old_video_tags_format(output_directory):
"title": "{title}", "title": "{title}",
} }
}, },
"overrides": {"music_video_artist": "JMC"}, "overrides": {
"music_video_artist": "JMC",
"music_video_directory": output_directory,
},
} }
@ -43,7 +45,6 @@ def single_video_preset_dict(output_directory):
"download": "https://youtube.com/watch?v=HKTNxEqsN3Q", "download": "https://youtube.com/watch?v=HKTNxEqsN3Q",
# override the output directory with our fixture-generated dir # override the output directory with our fixture-generated dir
"output_options": { "output_options": {
"output_directory": output_directory,
"maintain_download_archive": False, "maintain_download_archive": False,
}, },
"embed_thumbnail": True, # embed thumb into the video "embed_thumbnail": True, # embed thumb into the video
@ -52,7 +53,12 @@ def single_video_preset_dict(output_directory):
"video_tags": { "video_tags": {
"title": "{title}", "title": "{title}",
}, },
"overrides": {"music_video_artist": "JMC"}, # And test subscription download proba = 1.0
"throttle_protection": {"subscription_download_probability": 1.0},
"overrides": {
"music_video_artist": "JMC",
"music_video_directory": output_directory,
},
} }
@ -113,7 +119,7 @@ class TestYoutubeVideo:
transaction_log_summary_file_name="youtube/test_video.txt", transaction_log_summary_file_name="youtube/test_video.txt",
) )
@pytest.mark.parametrize("dry_run", [True]) @pytest.mark.parametrize("dry_run", [True, False])
def test_single_video_download( def test_single_video_download(
self, self,
default_config, default_config,

View file

@ -1,6 +1,5 @@
{ {
"JMC/JMC - Oblivion Mod Falcor p.1-thumb.jpg": "fb95b510681676e81c321171fc23143e", "JMC/Oblivion Mod Falcor p.1.jpg": "fb95b510681676e81c321171fc23143e",
"JMC/JMC - Oblivion Mod Falcor p.1.info.json": "08b0f7d93488d625bd7311ab925cec77", "JMC/Oblivion Mod Falcor p.1.mp4": "0448c9fd3eeaba4eca7f650fb93fe21b",
"JMC/JMC - Oblivion Mod Falcor p.1.mp4": "797b44f3207be01651780d6d86cb70bb", "JMC/Oblivion Mod Falcor p.1.nfo": "58c2be339869b5d071c1758d55c72ddb"
"JMC/JMC - Oblivion Mod Falcor p.1.nfo": "24cc4e17d2bebc89b2759ce5471d403e"
} }

View file

@ -5,8 +5,8 @@ from typing import Optional
import pytest import pytest
from ytdl_sub.config.config_file import ConfigFile from ytdl_sub.config.config_file import ConfigFile
from ytdl_sub.config.plugin_mapping import PluginMapping
from ytdl_sub.config.preset import PRESET_KEYS from ytdl_sub.config.preset import PRESET_KEYS
from ytdl_sub.config.preset_class_mappings import PluginMapping
from ytdl_sub.utils.exceptions import ValidationException from ytdl_sub.utils.exceptions import ValidationException

View file

@ -0,0 +1,59 @@
from conftest import assert_logs
from ytdl_sub.plugins.throttle_protection import logger as throttle_protection_logger
from ytdl_sub.subscriptions.subscription import Subscription
class TestThrottleProtectionPlugin:
def test_sleeps_log(
self,
config,
subscription_name,
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,
)
with mock_download_collection_entries(
is_youtube_channel=False, num_urls=1, is_extracted_audio=False
), assert_logs(
logger=throttle_protection_logger,
expected_message="Sleeping between downloads for %0.2f seconds",
log_level="debug",
expected_occurrences=4,
):
_ = subscription.download(dry_run=False)
with mock_download_collection_entries(
is_youtube_channel=False, num_urls=1, is_extracted_audio=False
), assert_logs(
logger=throttle_protection_logger,
expected_message="Sleeping between subscriptions for %0.2f seconds",
log_level="debug",
expected_occurrences=1,
):
_ = subscription.download(dry_run=False)