danvics/ytdl-sub
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix mocks

Browse source
This commit is contained in:
Jesse Bannon 2023-03-13 14:11:39 -07:00
parent 6a3e9cddc0
commit dd11e293f4
8 changed files with 216 additions and 205 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
src/ytdl_sub/downloaders/downloader.py
View file

@ -20,8 +20,7 @@ from ytdl_sub.downloaders.generic.validators import MultiUrlValidator
from ytdl_sub.downloaders.generic.validators import UrlThumbnailListValidator
from ytdl_sub.downloaders.generic.validators import UrlValidator
from ytdl_sub.downloaders.ytdl_options_builder import YTDLOptionsBuilder
from ytdl_sub.downloaders.ytdlp import extract_info_via_info_json
from ytdl_sub.downloaders.ytdlp import extract_info_with_retry
from ytdl_sub.downloaders.ytdlp import YTDLP
from ytdl_sub.entries.entry import Entry
from ytdl_sub.entries.entry_parent import EntryParent
from ytdl_sub.entries.variables.kwargs import COMMENTS
@ -266,7 +265,7 @@ class YtDlpDownloader(BaseDownloader[DownloaderOptionsT], ABC):
FileHandler.delete(info_json_file)
def _extract_entry_info_with_retry(self, entry: Entry) -> Entry:
download_entry_dict = extract_info_with_retry(
download_entry_dict = YTDLP.extract_info_with_retry(
ytdl_options_overrides=self.download_ytdl_options,
is_downloaded_fn=None if self.is_dry_run else entry.is_downloaded,
is_thumbnail_downloaded_fn=None
@ -334,7 +333,7 @@ class YtDlpDownloader(BaseDownloader[DownloaderOptionsT], ABC):
url = self.overrides.apply_formatter(collection_url.url)
with self._separate_download_archives():
entry_dicts = extract_info_via_info_json(
entry_dicts = YTDLP.extract_info_via_info_json(
working_directory=self.working_directory,
ytdl_options_overrides=self.metadata_ytdl_options,
log_prefix_on_info_json_dl="Downloading metadata for",

362
src/ytdl_sub/downloaders/ytdlp.py
View file

@ -19,199 +19,209 @@ from ytdl_sub.thread.log_entries_downloaded_listener import LogEntriesDownloaded
from ytdl_sub.utils.exceptions import FileNotDownloadedException
from ytdl_sub.utils.logger import Logger
download_logger = Logger.get(name="yt-dlp-downloader")
_extract_entry_num_retries: int = 5
_extract_entry_retry_wait_sec: int = 5
class YTDLP:
_EXTRACT_ENTRY_NUM_RETRIES: int = 5
_EXTRACT_ENTRY_RETRY_WAIT_SEC: int = 5
logger = Logger.get(name="yt-dlp-downloader")
@contextmanager
def ytdlp_downloader(ytdl_options_overrides: Dict) -> ytdl.YoutubeDL:
"""
Context manager to interact with yt_dlp.
"""
download_logger.debug("ytdl_options: %s", str(ytdl_options_overrides))
with Logger.handle_external_logs(name="yt-dlp"):
# Deep copy ytdl_options in case yt-dlp modifies the dict
with ytdl.YoutubeDL(copy.deepcopy(ytdl_options_overrides)) as ytdl_downloader:
yield ytdl_downloader
@classmethod
@contextmanager
def ytdlp_downloader(cls, ytdl_options_overrides: Dict) -> ytdl.YoutubeDL:
"""
Context manager to interact with yt_dlp.
"""
cls.logger.debug("ytdl_options: %s", str(ytdl_options_overrides))
with Logger.handle_external_logs(name="yt-dlp"):
# Deep copy ytdl_options in case yt-dlp modifies the dict
with ytdl.YoutubeDL(copy.deepcopy(ytdl_options_overrides)) as ytdl_downloader:
yield ytdl_downloader
@classmethod
def extract_info(cls, ytdl_options_overrides: Dict, **kwargs) -> Dict:
"""
Wrapper around yt_dlp.YoutubeDL.YoutubeDL.extract_info
All kwargs will passed to the extract_info function.
def extract_info(ytdl_options_overrides: Dict, **kwargs) -> Dict:
"""
Wrapper around yt_dlp.YoutubeDL.YoutubeDL.extract_info
All kwargs will passed to the extract_info function.
Parameters
----------
ytdl_options_overrides
Optional. Dict containing ytdl args to override other predefined ytdl args
**kwargs
arguments passed directory to YoutubeDL extract_info
"""
with cls.ytdlp_downloader(ytdl_options_overrides) as ytdlp:
return ytdlp.extract_info(**kwargs)
Parameters
----------
ytdl_options_overrides
Optional. Dict containing ytdl args to override other predefined ytdl args
**kwargs
arguments passed directory to YoutubeDL extract_info
"""
with ytdlp_downloader(ytdl_options_overrides) as ytdlp:
return ytdlp.extract_info(**kwargs)
@classmethod
def extract_info_with_retry(
cls,
ytdl_options_overrides: Dict,
is_downloaded_fn: Optional[Callable[[], bool]] = None,
is_thumbnail_downloaded_fn: Optional[Callable[[], bool]] = None,
**kwargs,
) -> Dict:
"""
Wrapper around yt_dlp.YoutubeDL.YoutubeDL.extract_info
All kwargs will passed to the extract_info function.
This should be used when downloading a single entry. Checks if the entry's video
and thumbnail files exist - retry if they do not.
def extract_info_with_retry(
ytdl_options_overrides: Dict,
is_downloaded_fn: Optional[Callable[[], bool]] = None,
is_thumbnail_downloaded_fn: Optional[Callable[[], bool]] = None,
**kwargs,
) -> Dict:
"""
Wrapper around yt_dlp.YoutubeDL.YoutubeDL.extract_info
All kwargs will passed to the extract_info function.
Parameters
----------
ytdl_options_overrides
Dict containing ytdl args to override other predefined ytdl args
is_downloaded_fn
Optional. Function to check if the entry is downloaded
is_thumbnail_downloaded_fn
Optional. Function to check if the entry thumbnail is downloaded
**kwargs
arguments passed directory to YoutubeDL extract_info
This should be used when downloading a single entry. Checks if the entry's video
and thumbnail files exist - retry if they do not.
Raises
------
FileNotDownloadedException
If the entry fails to download
"""
num_tries = 0
entry_files_exist = False
copied_ytdl_options_overrides = copy.deepcopy(ytdl_options_overrides)
Parameters
----------
ytdl_options_overrides
Dict containing ytdl args to override other predefined ytdl args
is_downloaded_fn
Optional. Function to check if the entry is downloaded
is_thumbnail_downloaded_fn
Optional. Function to check if the entry thumbnail is downloaded
**kwargs
arguments passed directory to YoutubeDL extract_info
Raises
------
FileNotDownloadedException
If the entry fails to download
"""
num_tries = 0
entry_files_exist = False
copied_ytdl_options_overrides = copy.deepcopy(ytdl_options_overrides)
while not entry_files_exist and num_tries < _extract_entry_num_retries:
entry_dict = extract_info(ytdl_options_overrides=copied_ytdl_options_overrides, **kwargs)
is_downloaded = is_downloaded_fn is None or is_downloaded_fn()
is_thumbnail_downloaded = is_thumbnail_downloaded_fn is None or is_thumbnail_downloaded_fn()
if is_downloaded and is_thumbnail_downloaded:
return entry_dict
# If the video file is downloaded but the thumbnail is not, then do not download
# the video again
if is_downloaded and not is_thumbnail_downloaded:
copied_ytdl_options_overrides["skip_download"] = True
copied_ytdl_options_overrides["writethumbnail"] = True
time.sleep(_extract_entry_retry_wait_sec)
num_tries += 1
# Remove the download archive so it can retry without thinking its already downloaded,
# even though it is not
if "download_archive" in copied_ytdl_options_overrides:
del copied_ytdl_options_overrides["download_archive"]
if num_tries < _extract_entry_num_retries:
download_logger.debug(
"Failed to download entry. Retrying %d / %d",
num_tries,
_extract_entry_num_retries,
while not entry_files_exist and num_tries < cls._EXTRACT_ENTRY_NUM_RETRIES:
entry_dict = cls.extract_info(
ytdl_options_overrides=copied_ytdl_options_overrides, **kwargs
)
error_dict = {"ytdl_options": ytdl_options_overrides, "kwargs": kwargs}
raise FileNotDownloadedException(
f"yt-dlp failed to download an entry with these arguments: {error_dict}"
)
is_downloaded = is_downloaded_fn is None or is_downloaded_fn()
is_thumbnail_downloaded = (
is_thumbnail_downloaded_fn is None or is_thumbnail_downloaded_fn()
)
if is_downloaded and is_thumbnail_downloaded:
return entry_dict
def _get_entry_dicts_from_info_json_files(working_directory: str) -> List[Dict]:
"""
Parameters
----------
working_directory
Directory that info json files are located
# If the video file is downloaded but the thumbnail is not, then do not download
# the video again
if is_downloaded and not is_thumbnail_downloaded:
copied_ytdl_options_overrides["skip_download"] = True
copied_ytdl_options_overrides["writethumbnail"] = True
Returns
-------
List of all info.json files read as JSON dicts
"""
entry_dicts: List[Dict] = []
info_json_paths = [
Path(working_directory) / file_name
for file_name in os.listdir(working_directory)
if file_name.endswith(".info.json")
]
time.sleep(cls._EXTRACT_ENTRY_RETRY_WAIT_SEC)
num_tries += 1
for info_json_path in info_json_paths:
with open(info_json_path, "r", encoding="utf-8") as file:
entry_dicts.append(json.load(file))
# Remove the download archive so it can retry without thinking its already downloaded,
# even though it is not
if "download_archive" in copied_ytdl_options_overrides:
del copied_ytdl_options_overrides["download_archive"]
return entry_dicts
if num_tries < cls._EXTRACT_ENTRY_NUM_RETRIES:
cls.logger.debug(
"Failed to download entry. Retrying %d / %d",
num_tries,
cls._EXTRACT_ENTRY_NUM_RETRIES,
)
@contextlib.contextmanager
def _listen_and_log_downloaded_info_json(working_directory: str, log_prefix: Optional[str]):
"""
Context manager that starts a separate thread that listens for new .info.json files,
prints their titles as they appear
"""
if not log_prefix:
yield
return
info_json_listener = LogEntriesDownloadedListener(
working_directory=working_directory,
log_prefix=log_prefix,
)
info_json_listener.start()
try:
yield
finally:
info_json_listener.complete = True
def extract_info_via_info_json(
working_directory: str,
ytdl_options_overrides: Dict,
log_prefix_on_info_json_dl: Optional[str] = None,
**kwargs,
) -> List[Dict]:
"""
Wrapper around yt_dlp.YoutubeDL.YoutubeDL.extract_info with infojson enabled. Entry dicts
are extracted via reading all info.json files in the working directory rather than
from the output of extract_info.
This allows us to catch RejectedVideoReached and ExistingVideoReached exceptions, and
simply ignore while still being able to read downloaded entry metadata.
Parameters
----------
working_directory
Directory that info json files reside in
ytdl_options_overrides
Dict containing ytdl args to override other predefined ytdl args
log_prefix_on_info_json_dl
Optional. Spin a new thread to listen for new info.json files. Log
f'{log_prefix_on_info_json_dl} {title}' when a new one appears
**kwargs
arguments passed directory to YoutubeDL extract_info
"""
try:
with _listen_and_log_downloaded_info_json(
working_directory=working_directory, log_prefix=log_prefix_on_info_json_dl
):
_ = extract_info(ytdl_options_overrides=ytdl_options_overrides, **kwargs)
except RejectedVideoReached:
download_logger.debug(
"RejectedVideoReached, stopping additional downloads "
"(Can be disable by setting `ytdl_options.break_on_reject` to False)."
error_dict = {"ytdl_options": ytdl_options_overrides, "kwargs": kwargs}
raise FileNotDownloadedException(
f"yt-dlp failed to download an entry with these arguments: {error_dict}"
)
except ExistingVideoReached:
download_logger.debug(
"ExistingVideoReached, stopping additional downloads. "
"(Can be disable by setting `ytdl_options.break_on_existing` to False)."
)
except MaxDownloadsReached:
download_logger.info("MaxDownloadsReached, stopping additional downloads.")
return _get_entry_dicts_from_info_json_files(working_directory=working_directory)
@classmethod
def _get_entry_dicts_from_info_json_files(cls, working_directory: str) -> List[Dict]:
"""
Parameters
----------
working_directory
Directory that info json files are located
Returns
-------
List of all info.json files read as JSON dicts
"""
entry_dicts: List[Dict] = []
info_json_paths = [
Path(working_directory) / file_name
for file_name in os.listdir(working_directory)
if file_name.endswith(".info.json")
]
for info_json_path in info_json_paths:
with open(info_json_path, "r", encoding="utf-8") as file:
entry_dicts.append(json.load(file))
return entry_dicts
@classmethod
@contextlib.contextmanager
def _listen_and_log_downloaded_info_json(
cls, working_directory: str, log_prefix: Optional[str]
):
"""
Context manager that starts a separate thread that listens for new .info.json files,
prints their titles as they appear
"""
if not log_prefix:
yield
return
info_json_listener = LogEntriesDownloadedListener(
working_directory=working_directory,
log_prefix=log_prefix,
)
info_json_listener.start()
try:
yield
finally:
info_json_listener.complete = True
@classmethod
def extract_info_via_info_json(
cls,
working_directory: str,
ytdl_options_overrides: Dict,
log_prefix_on_info_json_dl: Optional[str] = None,
**kwargs,
) -> List[Dict]:
"""
Wrapper around yt_dlp.YoutubeDL.YoutubeDL.extract_info with infojson enabled. Entry dicts
are extracted via reading all info.json files in the working directory rather than
from the output of extract_info.
This allows us to catch RejectedVideoReached and ExistingVideoReached exceptions, and
simply ignore while still being able to read downloaded entry metadata.
Parameters
----------
working_directory
Directory that info json files reside in
ytdl_options_overrides
Dict containing ytdl args to override other predefined ytdl args
log_prefix_on_info_json_dl
Optional. Spin a new thread to listen for new info.json files. Log
f'{log_prefix_on_info_json_dl} {title}' when a new one appears
**kwargs
arguments passed directory to YoutubeDL extract_info
"""
try:
with cls._listen_and_log_downloaded_info_json(
working_directory=working_directory, log_prefix=log_prefix_on_info_json_dl
):
_ = cls.extract_info(ytdl_options_overrides=ytdl_options_overrides, **kwargs)
except RejectedVideoReached:
cls.logger.debug(
"RejectedVideoReached, stopping additional downloads "
"(Can be disable by setting `ytdl_options.break_on_reject` to False)."
)
except ExistingVideoReached:
cls.logger.debug(
"ExistingVideoReached, stopping additional downloads. "
"(Can be disable by setting `ytdl_options.break_on_existing` to False)."
)
except MaxDownloadsReached:
cls.logger.info("MaxDownloadsReached, stopping additional downloads.")
return cls._get_entry_dicts_from_info_json_files(working_directory=working_directory)

3
tests/e2e/bandcamp/test_bandcamp.py
View file

@ -4,6 +4,7 @@ from expected_download import assert_expected_downloads
from expected_transaction_log import assert_transaction_log_matches
import ytdl_sub.downloaders.downloader
from ytdl_sub.downloaders.ytdlp import YTDLP
from ytdl_sub.subscriptions.subscription import Subscription
@ -65,7 +66,7 @@ class TestBandcamp:
# Ensure another invocation will hit ExistingVideoReached
if not dry_run:
with assert_logs(
logger=ytdl_sub.downloaders.downloader.download_logger,
logger=YTDLP.logger,
expected_message="ExistingVideoReached, stopping additional downloads",
log_level="debug",
):

9
tests/e2e/plugins/test_date_range.py
View file

@ -7,6 +7,7 @@ from expected_download import assert_expected_downloads
from expected_transaction_log import assert_transaction_log_matches
import ytdl_sub.downloaders.downloader
from ytdl_sub.downloaders.ytdlp import YTDLP
from ytdl_sub.subscriptions.subscription import Subscription
@ -68,7 +69,7 @@ class TestDateRange:
if not dry_run:
# try downloading again, ensure nothing more was downloaded
with assert_logs(
logger=ytdl_sub.downloaders.downloader.download_logger,
logger=YTDLP.logger,
expected_message="ExistingVideoReached, stopping additional downloads",
log_level="debug",
):
@ -139,7 +140,7 @@ class TestDateRange:
# First, download recent vids. Always download since we want to test dry-run
# on the rolling recent portion.
with assert_logs(
logger=ytdl_sub.downloaders.downloader.download_logger,
logger=YTDLP.logger,
expected_message="RejectedVideoReached, stopping additional downloads",
log_level="debug",
):
@ -159,7 +160,7 @@ class TestDateRange:
# Then, download the rolling recent vids subscription. This should remove one of the
# two videos
with assert_logs(
logger=ytdl_sub.downloaders.downloader.download_logger,
logger=YTDLP.logger,
expected_message="ExistingVideoReached, stopping additional downloads",
log_level="debug",
):
@ -186,7 +187,7 @@ class TestDateRange:
# existing
if not dry_run:
with assert_logs(
logger=ytdl_sub.downloaders.downloader.download_logger,
logger=YTDLP.logger,
expected_message="ExistingVideoReached, stopping additional downloads",
log_level="debug",
):

5
tests/e2e/youtube/test_playlist.py
View file

@ -5,6 +5,7 @@ from expected_download import assert_expected_downloads
from expected_transaction_log import assert_transaction_log_matches
import ytdl_sub.downloaders.downloader
from ytdl_sub.downloaders.ytdlp import YTDLP
from ytdl_sub.subscriptions.subscription import Subscription
@ -83,7 +84,7 @@ class TestPlaylist:
# Ensure another invocation will hit ExistingVideoReached
if not dry_run:
with assert_logs(
logger=ytdl_sub.downloaders.downloader.download_logger,
logger=YTDLP.logger,
expected_message="ExistingVideoReached, stopping additional downloads",
log_level="debug",
):
@ -131,7 +132,7 @@ class TestPlaylist:
if not dry_run:
# Ensure another invocation will hit ExistingVideoReached
with assert_logs(
logger=ytdl_sub.downloaders.downloader.download_logger,
logger=YTDLP.logger,
expected_message="ExistingVideoReached, stopping additional downloads",
log_level="debug",
):

1
tests/expected_download.py
View file

@ -1,6 +1,5 @@
import json
import os.path
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import List

30
tests/unit/prebuilt_presets/conftest.py
View file

@ -12,6 +12,7 @@ from resources import copy_file_fixture
from ytdl_sub.config.config_file import ConfigFile
from ytdl_sub.downloaders.downloader import YtDlpDownloader
from ytdl_sub.downloaders.ytdlp import YTDLP
from ytdl_sub.entries.variables.kwargs import DESCRIPTION
from ytdl_sub.entries.variables.kwargs import EPOCH
from ytdl_sub.entries.variables.kwargs import EXT
@ -100,23 +101,22 @@ def mock_entry_dict_factory(mock_downloaded_file_path) -> Callable:
@pytest.fixture
def mock_download_collection_thumbnail(mock_downloaded_file_path):
def _mock_download_thumbnail(output_path: str) -> bool:
# mock_file_factory(file_name=output_path.split("/")[-1])
output_name = os.path.basename(output_path)
def _mock_download_and_convert_url_thumbnail(
thumbnail_url: str, output_thumbnail_path: str
) -> bool:
_ = thumbnail_url
output_name = os.path.basename(output_thumbnail_path)
if "poster" in output_name or "show" in output_name:
copy_file_fixture(fixture_name="poster.jpg", output_file_path=output_path)
copy_file_fixture(fixture_name="poster.jpg", output_file_path=output_thumbnail_path)
return True
elif "fanart" in output_name:
copy_file_fixture(fixture_name="fanart.jpeg", output_file_path=output_path)
copy_file_fixture(fixture_name="fanart.jpeg", output_file_path=output_thumbnail_path)
return True
return False
with patch.object(
YtDlpDownloader,
"_download_thumbnail",
new=lambda _, thumbnail_url, output_thumbnail_path: _mock_download_thumbnail(
output_thumbnail_path
),
with patch(
"ytdl_sub.downloaders.downloader.download_and_convert_url_thumbnail",
new=_mock_download_and_convert_url_thumbnail,
):
yield # TODO: create file here
@ -126,11 +126,9 @@ def mock_download_collection_entries(
mock_download_collection_thumbnail, mock_entry_dict_factory: Callable, working_directory: str
):
@contextlib.contextmanager
def _mock_download_collection_entries_factory(is_youtube_channel: bool):
def _mock_download_collection_entries_factory(is_youtube_channel: bool, num_urls: int = 1):
def _write_entries_to_working_dir(*args, **kwargs) -> List[Dict]:
if (len(args[0].collection.urls.list) == 1) or (
"season.2" in kwargs["url"] and len(args[0].download_options.urls.list) > 1
):
if num_urls == 1 or ("season.2" in kwargs["url"] and num_urls > 1):
return [
mock_entry_dict_factory(
uid="21-1",
@ -202,7 +200,7 @@ def mock_download_collection_entries(
]
with patch.object(
YtDlpDownloader, "extract_info_via_info_json", new=_write_entries_to_working_dir
YTDLP, "extract_info_via_info_json", new=_write_entries_to_working_dir
), patch.object(
YtDlpDownloader, "_extract_entry_info_with_retry", new=lambda _, entry: entry
):

4
tests/unit/prebuilt_presets/test_prebuilt_presets.py
View file

@ -215,7 +215,9 @@ class TestPrebuiltTvShowCollectionPresets:
},
)
with mock_download_collection_entries(is_youtube_channel=is_youtube_channel):
with mock_download_collection_entries(
is_youtube_channel=is_youtube_channel, num_urls=len(season_indices)
):
transaction_log = subscription.download(dry_run=False)
assert_transaction_log_matches(