[DEV] Partially unit test throttle protection (#1065)

This commit is contained in:
Jesse Bannon 2024-09-26 10:27:31 -07:00 committed by GitHub
parent 79ba60021c
commit 513d1f5a97
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 55 additions and 39 deletions

View file

@ -1,24 +1,8 @@
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 {
@ -34,27 +18,6 @@ def preset_dict_subscription_download_proba_0(output_directory):
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,

View file

@ -1,5 +1,6 @@
import contextlib
import os
import shutil
from pathlib import Path
from typing import Callable
from typing import Dict
@ -133,8 +134,13 @@ def mock_download_collection_entries(
):
@contextlib.contextmanager
def _mock_download_collection_entries_factory(
is_youtube_channel: bool, num_urls: int = 1, is_extracted_audio: bool = False
is_youtube_channel: bool,
num_urls: int = 1,
is_extracted_audio: bool = False,
is_dry_run: bool = False,
):
is_real_run = not is_dry_run
def _write_entries_to_working_dir(*args, **kwargs) -> List[Dict]:
# Second TV URL or second soundcloud URL, which downloads first
is_second_url = "2" in kwargs["url"] or kwargs["url"].endswith("/albums")
@ -149,6 +155,7 @@ def mock_download_collection_entries(
playlist_count=4,
is_youtube_channel=is_youtube_channel,
is_extracted_audio=is_extracted_audio,
mock_download_to_working_dir=is_real_run,
), # 1
mock_entry_dict_factory(
uid="20-1",
@ -158,6 +165,7 @@ def mock_download_collection_entries(
playlist_count=4,
is_youtube_channel=is_youtube_channel,
is_extracted_audio=is_extracted_audio,
mock_download_to_working_dir=is_real_run,
), # 2 98
mock_entry_dict_factory(
uid="20-2",
@ -167,6 +175,7 @@ def mock_download_collection_entries(
playlist_count=4,
is_youtube_channel=is_youtube_channel,
is_extracted_audio=is_extracted_audio,
mock_download_to_working_dir=is_real_run,
), # 1 99
mock_entry_dict_factory(
uid="20-3",
@ -176,6 +185,7 @@ def mock_download_collection_entries(
playlist_count=4,
is_youtube_channel=is_youtube_channel,
is_extracted_audio=is_extracted_audio,
mock_download_to_working_dir=is_real_run,
),
]
return [
@ -187,8 +197,8 @@ def mock_download_collection_entries(
playlist_index=1,
playlist_count=5,
is_youtube_channel=is_youtube_channel,
mock_download_to_working_dir=False,
is_extracted_audio=is_extracted_audio,
mock_download_to_working_dir=False,
),
mock_entry_dict_factory(
uid="20-4",
@ -198,6 +208,7 @@ def mock_download_collection_entries(
playlist_count=5,
is_youtube_channel=is_youtube_channel,
is_extracted_audio=is_extracted_audio,
mock_download_to_working_dir=is_real_run,
),
mock_entry_dict_factory(
uid="20-5",
@ -207,6 +218,7 @@ def mock_download_collection_entries(
playlist_count=5,
is_youtube_channel=is_youtube_channel,
is_extracted_audio=is_extracted_audio,
mock_download_to_working_dir=is_real_run,
),
mock_entry_dict_factory(
uid="20-6",
@ -216,6 +228,7 @@ def mock_download_collection_entries(
playlist_count=5,
is_youtube_channel=is_youtube_channel,
is_extracted_audio=is_extracted_audio,
mock_download_to_working_dir=is_real_run,
),
mock_entry_dict_factory(
uid="20-7",
@ -225,6 +238,7 @@ def mock_download_collection_entries(
playlist_count=5,
is_youtube_channel=is_youtube_channel,
is_extracted_audio=is_extracted_audio,
mock_download_to_working_dir=is_real_run,
),
]

View file

@ -123,3 +123,42 @@ class TestThrottleProtectionPlugin:
),
):
_ = subscription.download(dry_run=False)
def test_max_downloads(
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": {"max_downloads_per_subscription": {"max": 0}},
}
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, is_dry_run=True
),
assert_logs(
logger=throttle_protection_logger,
expected_message="Reached subscription max downloads of %d",
log_level="info",
expected_occurrences=1,
),
):
transaction_log = subscription.download(dry_run=True)
assert transaction_log.is_empty