diff --git a/src/ytdl_sub/subscriptions/subscription_download.py b/src/ytdl_sub/subscriptions/subscription_download.py index 8b20c91e..9915a18f 100644 --- a/src/ytdl_sub/subscriptions/subscription_download.py +++ b/src/ytdl_sub/subscriptions/subscription_download.py @@ -262,6 +262,31 @@ class SubscriptionDownload(BaseSubscription, ABC): self._cleanup_entry_files(entry) + def _process_subscription( + self, + plugins: List[Plugin], + entries: Iterable[Entry] | Iterable[Tuple[Entry, FileMetadata]], + dry_run: bool, + ) -> FileHandlerTransactionLog: + for entry in entries: + entry_metadata = FileMetadata() + if isinstance(entry, tuple): + entry, entry_metadata = entry + + if split_plugin := _get_split_plugin(plugins): + self._process_split_entry( + split_plugin=split_plugin, plugins=plugins, dry_run=dry_run, entry=entry + ) + else: + self._process_entry( + plugins=plugins, dry_run=dry_run, entry=entry, entry_metadata=entry_metadata + ) + + for plugin in plugins: + plugin.post_process_subscription() + + return self._enhanced_download_archive.get_file_handler_transaction_log() + def download(self, dry_run: bool = False) -> FileHandlerTransactionLog: """ Performs the subscription download @@ -292,24 +317,11 @@ class SubscriptionDownload(BaseSubscription, ABC): overrides=self.overrides, ) - for entry in downloader.download(): - entry_metadata = FileMetadata() - if isinstance(entry, tuple): - entry, entry_metadata = entry - - if split_plugin := _get_split_plugin(plugins): - self._process_split_entry( - split_plugin=split_plugin, plugins=plugins, dry_run=dry_run, entry=entry - ) - else: - self._process_entry( - plugins=plugins, dry_run=dry_run, entry=entry, entry_metadata=entry_metadata - ) - - for plugin in plugins: - plugin.post_process_subscription() - - return self._enhanced_download_archive.get_file_handler_transaction_log() + return self._process_subscription( + plugins=plugins, + entries=downloader.download(), + dry_run=dry_run, + ) def _get_entries_for_reformat( self, original_enhanced_download_archive: EnhancedDownloadArchive @@ -328,10 +340,10 @@ class SubscriptionDownload(BaseSubscription, ABC): entry_dict=json.load(maybe_info_json), working_directory=self.working_directory, ) - except Exception: + except Exception as exc: raise ValidationException( "info.json file cannot be loaded - subscription cannot be reformatted" - ) + ) from exc if not maybe_entry: raise ValidationException( @@ -356,14 +368,12 @@ class SubscriptionDownload(BaseSubscription, ABC): yield entry - def reformat( - self, reformat_output_directory: Path, dry_run: bool = False - ) -> FileHandlerTransactionLog: + def reformat(self, dry_run: bool = False) -> FileHandlerTransactionLog: original_enhanced_download_archive = self._enhanced_download_archive self._enhanced_download_archive = EnhancedDownloadArchive( subscription_name=self.name, working_directory=self.working_directory, - output_directory=reformat_output_directory, + output_directory=self.output_directory, dry_run=dry_run, ) @@ -371,23 +381,10 @@ class SubscriptionDownload(BaseSubscription, ABC): plugins = self._initialize_plugins() with self._subscription_download_context_managers(): - for entry in self._get_entries_for_reformat( - original_enhanced_download_archive=original_enhanced_download_archive - ): - entry_metadata = FileMetadata() - if isinstance(entry, tuple): - entry, entry_metadata = entry - - if split_plugin := _get_split_plugin(plugins): - self._process_split_entry( - split_plugin=split_plugin, plugins=plugins, dry_run=dry_run, entry=entry - ) - else: - self._process_entry( - plugins=plugins, dry_run=dry_run, entry=entry, entry_metadata=entry_metadata - ) - - for plugin in plugins: - plugin.post_process_subscription() - - return self._enhanced_download_archive.get_file_handler_transaction_log() + return self._process_subscription( + plugins=plugins, + entries=self._get_entries_for_reformat( + original_enhanced_download_archive=original_enhanced_download_archive + ), + dry_run=dry_run, + ) diff --git a/tests/conftest.py b/tests/conftest.py index 66abf88a..54fab744 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -70,11 +70,13 @@ def output_directory() -> Path: with tempfile.TemporaryDirectory() as temp_dir: yield temp_dir + @pytest.fixture() def reformat_directory() -> Path: with tempfile.TemporaryDirectory() as temp_dir: yield temp_dir + @contextlib.contextmanager def assert_logs(logger: logging.Logger, expected_message: str, log_level: str = "debug"): """ diff --git a/tests/e2e/youtube/test_channel.py b/tests/e2e/youtube/test_channel.py index 2b179076..97ff436e 100644 --- a/tests/e2e/youtube/test_channel.py +++ b/tests/e2e/youtube/test_channel.py @@ -1,10 +1,9 @@ from pathlib import Path import pytest -from yt_dlp.utils import sanitize_filename - from expected_download import assert_expected_downloads from expected_transaction_log import assert_transaction_log_matches +from yt_dlp.utils import sanitize_filename from ytdl_sub.subscriptions.subscription import Subscription @@ -80,16 +79,14 @@ class TestChannel: if not reformat: return - reformat_directory = Path(reformat_directory) / sanitize_filename("Project / Zombie") - full_channel_subscription.reformat(reformat_output_directory=reformat_directory, dry_run=dry_run) + full_channel_subscription.reformat(dry_run=dry_run) assert_transaction_log_matches( - output_directory=reformat_directory, + output_directory=output_directory, transaction_log=transaction_log, transaction_log_summary_file_name="youtube/test_channel_full.txt", ) assert_expected_downloads( - output_directory=reformat_directory, + output_directory=output_directory, dry_run=dry_run, expected_download_summary_file_name="youtube/test_channel_full.json", ) -