diff --git a/src/ytdl_sub/subscriptions/subscription_download.py b/src/ytdl_sub/subscriptions/subscription_download.py index 7cdc2144..05460ba1 100644 --- a/src/ytdl_sub/subscriptions/subscription_download.py +++ b/src/ytdl_sub/subscriptions/subscription_download.py @@ -95,9 +95,10 @@ class SubscriptionDownload(BaseSubscription, ABC): entry=entry, ) - def _delete_working_directory(self, is_error: bool) -> None: + def _delete_working_directory(self, is_error: bool = False) -> None: _ = is_error - shutil.rmtree(self.working_directory) + if os.path.isdir(self.working_directory): + shutil.rmtree(self.working_directory) @contextlib.contextmanager def _prepare_working_directory(self): @@ -105,6 +106,7 @@ class SubscriptionDownload(BaseSubscription, ABC): Context manager to create all directories to the working directory. Deletes the entire working directory when cleaning up. """ + self._delete_working_directory() os.makedirs(self.working_directory, exist_ok=True) try: @@ -113,7 +115,7 @@ class SubscriptionDownload(BaseSubscription, ABC): self._delete_working_directory(is_error=True) raise exc else: - self._delete_working_directory(is_error=False) + self._delete_working_directory() @contextlib.contextmanager def _maintain_archive_file(self): diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index de447ec1..aadbb1d8 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -1,4 +1,5 @@ import json +import os import shutil import sys import tempfile @@ -27,10 +28,11 @@ def working_directory() -> str: """ with tempfile.TemporaryDirectory() as temp_dir: - def _assert_working_directory_empty(self, is_error: bool): + def _assert_working_directory_empty(self, is_error: bool = False): files = [str(file_path) for file_path in _get_files_in_directory(temp_dir)] num_files = len(files) - shutil.rmtree(temp_dir) + if os.path.isdir(temp_dir): + shutil.rmtree(temp_dir) if not is_error: if num_files > 0: diff --git a/tests/unit/prebuilt_presets/conftest.py b/tests/unit/prebuilt_presets/conftest.py index 7ce3d4e8..9401b3d2 100644 --- a/tests/unit/prebuilt_presets/conftest.py +++ b/tests/unit/prebuilt_presets/conftest.py @@ -4,6 +4,7 @@ import tempfile from pathlib import Path from typing import Callable from typing import Dict +from typing import List from unittest.mock import patch import pytest @@ -126,85 +127,83 @@ def mock_download_collection_entries( def _(**kwargs): return mock_entry_dict_factory(**kwargs) - collection_1_entry_dicts = [ - _( - uid="21-1", - upload_date="20210808", - playlist_index=1, - playlist_count=4, - is_youtube_channel=is_youtube_channel, - ), # 1 - _( - uid="20-1", - upload_date="20200808", - playlist_index=2, - playlist_count=4, - is_youtube_channel=is_youtube_channel, - ), # 2 98 - _( - uid="20-2", - upload_date="20200808", - playlist_index=3, - playlist_count=4, - is_youtube_channel=is_youtube_channel, - ), # 1 99 - _( - uid="20-3", - upload_date="20200807", - playlist_index=4, - playlist_count=4, - is_youtube_channel=is_youtube_channel, - ), - ] - collection_2_entry_dicts = [ - # 20-3 should resolve to collection 1 (which is season 2) - _( - uid="20-3", - upload_date="20200807", - playlist_index=1, - playlist_count=5, - is_youtube_channel=is_youtube_channel, - ), - _( - uid="20-4", - upload_date="20200806", - playlist_index=2, - playlist_count=5, - is_youtube_channel=is_youtube_channel, - ), - _( - uid="20-5", - upload_date="20200706", - playlist_index=3, - playlist_count=5, - is_youtube_channel=is_youtube_channel, - ), - _( - uid="20-6", - upload_date="20200706", - playlist_index=4, - playlist_count=5, - is_youtube_channel=is_youtube_channel, - ), - _( - uid="20-7", - upload_date="20200606", - playlist_index=5, - playlist_count=5, - is_youtube_channel=is_youtube_channel, - ), - ] + 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 + ): + return [ + _( + uid="21-1", + upload_date="20210808", + playlist_index=1, + playlist_count=4, + is_youtube_channel=is_youtube_channel, + ), # 1 + _( + uid="20-1", + upload_date="20200808", + playlist_index=2, + playlist_count=4, + is_youtube_channel=is_youtube_channel, + ), # 2 98 + _( + uid="20-2", + upload_date="20200808", + playlist_index=3, + playlist_count=4, + is_youtube_channel=is_youtube_channel, + ), # 1 99 + _( + uid="20-3", + upload_date="20200807", + playlist_index=4, + playlist_count=4, + is_youtube_channel=is_youtube_channel, + ), + ] + return [ + # 20-3 should resolve to collection 1 (which is season 2) + _( + uid="20-3", + upload_date="20200807", + playlist_index=1, + playlist_count=5, + is_youtube_channel=is_youtube_channel, + ), + _( + uid="20-4", + upload_date="20200806", + playlist_index=2, + playlist_count=5, + is_youtube_channel=is_youtube_channel, + ), + _( + uid="20-5", + upload_date="20200706", + playlist_index=3, + playlist_count=5, + is_youtube_channel=is_youtube_channel, + ), + _( + uid="20-6", + upload_date="20200706", + playlist_index=4, + playlist_count=5, + is_youtube_channel=is_youtube_channel, + ), + _( + uid="20-7", + upload_date="20200606", + playlist_index=5, + playlist_count=5, + is_youtube_channel=is_youtube_channel, + ), + ] with patch.object( - Downloader, "extract_info_via_info_json" - ) as mock_download_metadata, patch.object( - Downloader, "_extract_entry_info_with_retry", new=lambda _, entry: entry - ): + Downloader, "extract_info_via_info_json", new=_write_entries_to_working_dir + ), patch.object(Downloader, "_extract_entry_info_with_retry", new=lambda _, entry: entry): # Stub out metadata. TODO: update this if we do metadata plugins - mock_download_metadata.side_effect = [ - collection_1_entry_dicts, - collection_2_entry_dicts, - ] yield return _mock_download_collection_entries_factory