diff --git a/src/jobs/removal_handler.py b/src/jobs/removal_handler.py index 9244e01..05f7167 100644 --- a/src/jobs/removal_handler.py +++ b/src/jobs/removal_handler.py @@ -13,21 +13,21 @@ class RemovalHandler: str(self.arr.tracker.deleted), ) - queue_item = affected_downloads[download_id][0] - handling_method = await self._get_handling_method(download_id, queue_item) + affected_download = affected_downloads[download_id] + handling_method = await self._get_handling_method(download_id, affected_download) if download_id in self.arr.tracker.deleted or handling_method == "skip": del affected_downloads[download_id] continue if handling_method == "remove": - await self._remove_download(queue_item, blocklist) + await self._remove_download(affected_download, blocklist) elif handling_method == "tag_as_obsolete": - await self._tag_as_obsolete(queue_item, download_id) + await self._tag_as_obsolete(affected_download, download_id) # Print out detailed removal messages (if any) - if "removal_messages" in queue_item: - for msg in queue_item["removal_messages"]: + if "removal_messages" in affected_download: + for msg in affected_download["removal_messages"]: logger.info(msg) self.arr.tracker.deleted.append(download_id) @@ -38,22 +38,22 @@ class RemovalHandler: ) - async def _remove_download(self, queue_item, blocklist): - queue_id = queue_item["id"] - logger.info(f">>> Job '{self.job_name}' triggered removal: {queue_item['title']}") + async def _remove_download(self, affected_download, blocklist): + queue_id = affected_download["queue_ids"][0] + logger.info(f">>> Job '{self.job_name}' triggered removal: {affected_download['title']}") await self.arr.remove_queue_item(queue_id=queue_id, blocklist=blocklist) - async def _tag_as_obsolete(self, queue_item, download_id): - logger.info(f">>> Job'{self.job_name}' triggered obsolete-tagging: {queue_item['title']}") + async def _tag_as_obsolete(self, affected_download, download_id): + logger.info(f">>> Job'{self.job_name}' triggered obsolete-tagging: {affected_download['title']}") for qbit in self.settings.download_clients.qbittorrent: await qbit.set_tag(tags=[self.settings.general.obsolete_tag], hashes=[download_id]) - async def _get_handling_method(self, download_id, queue_item): - if queue_item['protocol'] != 'torrent': + async def _get_handling_method(self, download_id, affected_download): + if affected_download['protocol'] != 'torrent': return "remove" # handling is only implemented for torrent - client_implemenation = await self.arr.get_download_client_implementation(queue_item['downloadClient']) + client_implemenation = await self.arr.get_download_client_implementation(affected_download['downloadClient']) if client_implemenation != "QBittorrent": return "remove" # handling is only implemented for qbit diff --git a/src/jobs/remove_unmonitored.py b/src/jobs/remove_unmonitored.py index 5c8f09b..ad20f5e 100644 --- a/src/jobs/remove_unmonitored.py +++ b/src/jobs/remove_unmonitored.py @@ -1,5 +1,6 @@ from src.jobs.removal_job import RemovalJob + class RemoveUnmonitored(RemovalJob): queue_scope = "normal" blocklist = False @@ -9,8 +10,11 @@ class RemoveUnmonitored(RemovalJob): monitored_download_ids = [] for item in self.queue: detail_item_id = item["detail_item_id"] - if await self.arr.is_monitored(detail_item_id): + if detail_item_id is None or await self.arr.is_monitored(detail_item_id): + # When queue item has been matched to artist (for instance in lidarr) but not yet to the detail (eg. album), then detail key is logically missing. + # Thus we can't check if the item is monitored yet monitored_download_ids.append(item["downloadId"]) + # Second pass: Append queue items none that depends on download id is monitored affected_items = [] @@ -19,4 +23,4 @@ class RemoveUnmonitored(RemovalJob): affected_items.append( queue_item ) # One downloadID may be shared by multiple queue_items. Only removes it if ALL queueitems are unmonitored - return affected_items \ No newline at end of file + return affected_items diff --git a/src/utils/queue_manager.py b/src/utils/queue_manager.py index 0e0bff6..47209cd 100644 --- a/src/utils/queue_manager.py +++ b/src/utils/queue_manager.py @@ -155,56 +155,40 @@ class QueueManager: def format_queue(self, queue_items): if not queue_items: return "empty" + return self.group_by_download_id(queue_items) - formatted_dict = {} + def group_by_download_id(self, queue_items): + # Groups queue items by download ID and returns a dict where download ID is the key, + # and the value is a dict with a list of IDs and other selected metadata. + retain_keys = [ + "detail_item_id", + "title", + "size", + "sizeleft", + "downloadClient", + "protocol", + "status", + "trackedDownloadState", + "statusMessages", + "removal_messages", + ] + + grouped_dict = {} for queue_item in queue_items: download_id = queue_item.get("downloadId") item_id = queue_item.get("id") - if download_id in formatted_dict: - formatted_dict[download_id]["IDs"].append(item_id) + if download_id in grouped_dict: + grouped_dict[download_id]["queue_ids"].append(item_id) else: - formatted_dict[download_id] = { - "downloadId": download_id, - "downloadTitle": queue_item.get("title"), - "protocol": [queue_item.get("protocol")], - "status": [queue_item.get("status")], - "IDs": [item_id], + grouped_dict[download_id] = { + "queue_ids": [item_id], + **{ + key: queue_item[key] + for key in retain_keys + if key in queue_item + }, } - return list(formatted_dict.values()) - - def group_by_download_id(self, queue_items): - # Groups queue items by download ID and returns a dict where download ID is the key, and value is the list of queue items belonging to that downloadID - # Queue item is limited to certain keys - retain_keys = { - "id": None, - "detail_item_id": None, - "title": "Unknown", - "size": 0, - "sizeleft": 0, - "downloadClient": "Unknown", - "protocol": "Unknown", - "status": "Unknown", - "trackedDownloadState": "Unknown", - "statusMessages": [], - "removal_messages": [], - } - - grouped_dict = {} - - for queue_item in queue_items: - download_id = queue_item["downloadId"] - if download_id not in grouped_dict: - grouped_dict[download_id] = [] - - # Filter and add default values if keys are missing - filtered_item = { - key: queue_item.get(key, retain_keys.get(key, None)) - for key in retain_keys - } - - grouped_dict[download_id].append(filtered_item) - return grouped_dict diff --git a/tests/jobs/test_removal_handler.py b/tests/jobs/test_removal_handler.py index c0701c6..5ae7ea7 100644 --- a/tests/jobs/test_removal_handler.py +++ b/tests/jobs/test_removal_handler.py @@ -1,128 +1,46 @@ -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock import pytest from src.jobs.removal_handler import RemovalHandler -# ---------- Fixtures ---------- -@pytest.fixture(name="mock_logger") -def fixture_mock_logger(): - with patch("src.jobs.removal_handler.logger") as mock: - yield mock - - -@pytest.fixture(name="settings") -def fixture_settings(): - settings = AsyncMock() - settings.general.test_run = False - settings.general.obsolete_tag = "obsolete_tag" - settings.download_clients.qbittorrent = [AsyncMock()] - return settings - - -@pytest.fixture(name="arr") -def fixture_arr(): +@pytest.mark.parametrize( + "qbittorrent_configured, is_private, client_impl, protocol, expected", + [ + (True, True, "QBittorrent", "torrent", "private_handling"), + (True, False, "QBittorrent", "torrent", "public_handling"), + (False, True, "QBittorrent", "torrent", "remove"), + (False, False, "QBittorrent", "torrent", "remove"), + (True, False, "Transmission", "torrent", "remove"), # unsupported client + (True, False, "MyUseNetClient", "usenet", "remove"), # unsupported protocol + ] +) +@pytest.mark.asyncio +async def test_get_handling_method( + qbittorrent_configured, + is_private, + client_impl, + protocol, + expected, +): + # Mock arr arr = AsyncMock() - arr.api_url = "https://mock-api-url" - arr.api_key = "mock_api_key" - arr.tracker = AsyncMock() - arr.tracker.deleted = [] - arr.get_download_client_implementation.return_value = "QBittorrent" - return arr + arr.tracker.private = ["A"] if is_private else [] + arr.get_download_client_implementation.return_value = client_impl + # Mock settings + settings = MagicMock() + settings.download_clients.qbittorrent = ["dummy"] if qbittorrent_configured else [] + settings.general.private_tracker_handling = "private_handling" + settings.general.public_tracker_handling = "public_handling" -@pytest.fixture(name="affected_downloads") -def fixture_affected_downloads(): - return { - "AABBCC": [ - { - "id": 1, - "downloadId": "AABBCC", - "title": "My Series A - Season 1", - "size": 1000, - "sizeleft": 500, - "downloadClient": "qBittorrent", - "protocol": "torrent", - "status": "paused", - "trackedDownloadState": "downloading", - "statusMessages": [], - } - ] + handler = RemovalHandler(arr=arr, settings=settings, job_name="test") + + affected_download = { + "downloadClient": "qBittorrent", + "protocol": protocol, } + result = await handler._get_handling_method("A", affected_download) # pylint: disable=protected-access + assert result == expected -# ---------- Parametrized Test ---------- - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "protocol, qb_config, client_impl, is_private, pub_handling, priv_handling, expected", - [ - ("emule", [AsyncMock()], "MyDonkey", None, "remove", "remove", "remove"), - ("torrent", [], "QBittorrent", None, "remove", "remove", "remove"), - ("torrent", [AsyncMock()], "OtherClient", None, "remove", "remove", "remove"), - - ("torrent", [AsyncMock()], "QBittorrent", True, "remove", "remove", "remove"), - ("torrent", [AsyncMock()], "QBittorrent", True, "remove", "tag_as_obsolete", "tag_as_obsolete"), - ("torrent", [AsyncMock()], "QBittorrent", True, "remove", "skip", "skip"), - - ("torrent", [AsyncMock()], "QBittorrent", False, "remove", "remove", "remove"), - ("torrent", [AsyncMock()], "QBittorrent", False, "tag_as_obsolete", "remove", "tag_as_obsolete"), - ("torrent", [AsyncMock()], "QBittorrent", False, "skip", "remove", "skip"), - ], -) -async def test_remove_downloads( - protocol, - qb_config, - client_impl, - is_private, - pub_handling, - priv_handling, - expected, - arr, - settings, - affected_downloads, -): - # ---------- Arrange ---------- - download_id = "AABBCC" - item = affected_downloads[download_id][0] - - item["protocol"] = protocol - item["downloadClient"] = "qBittorrent" - - settings.download_clients.qbittorrent = qb_config - settings.general.public_tracker_handling = pub_handling - settings.general.private_tracker_handling = priv_handling - - arr.get_download_client_implementation.return_value = client_impl - arr.tracker.private = [download_id] if is_private else [] - arr.tracker.deleted = [] - - handler = RemovalHandler(arr=arr, settings=settings, job_name="Test Job") - - # ---------- Act ---------- - await handler.remove_downloads(affected_downloads, blocklist=True) - observed = await handler._get_handling_method(download_id, item) - - # ---------- Assert ---------- - assert observed == expected - - if expected == "remove": - arr.remove_queue_item.assert_awaited_once_with( - queue_id=item["id"], blocklist=True - ) - assert download_id in arr.tracker.deleted - - elif expected == "tag_as_obsolete": - if qb_config: - qb_config[0].set_tag.assert_awaited_once_with( - tags=[settings.general.obsolete_tag], - hashes=[download_id], - ) - assert download_id in arr.tracker.deleted - - elif expected == "skip": - assert download_id not in affected_downloads - assert download_id not in arr.tracker.deleted - - if expected != "tag_as_obsolete" and qb_config: - qb_config[0].set_tag.assert_not_awaited() diff --git a/tests/jobs/test_remove_bad_files.py b/tests/jobs/test_remove_bad_files.py index d30eb22..1d8528c 100644 --- a/tests/jobs/test_remove_bad_files.py +++ b/tests/jobs/test_remove_bad_files.py @@ -4,25 +4,11 @@ import pytest from src.jobs.remove_bad_files import RemoveBadFiles # Fixture for arr mock -@pytest.fixture(name="arr") -def fixture_arr(): - arr = AsyncMock() - arr.api_url = "https://mock-api-url" - arr.api_key = "mock_api_key" - arr.tracker = AsyncMock() - arr.tracker.extension_checked = [] - arr.get_download_client_implementation.return_value = "QBittorrent" - return arr - - -@pytest.fixture(name="qbit_client") -def fixture_qbit_client(): - qbit_client = AsyncMock() - return qbit_client - - @pytest.fixture(name="removal_job") -def fixture_removal_job(arr): +def fixture_removal_job(): + arr = AsyncMock() + arr.get_download_client_implementation.return_value = "QBittorrent" + removal_job = RemoveBadFiles(arr=arr, settings=MagicMock(), job_name="test") removal_job.arr = arr removal_job.job = MagicMock() @@ -186,10 +172,11 @@ def test_is_complete_partial(removal_job, file, is_incomplete_partial): ], ) @pytest.mark.asyncio -async def test_get_items_to_process(qbit_item, expected_processed, removal_job, arr): +async def test_get_items_to_process(qbit_item, expected_processed, removal_job): """Test the _get_items_to_process method of RemoveBadFiles class.""" # Mocking the tracker extension_checked to simulate which torrents have been checked - arr.tracker.extension_checked = {"checked-hash"} + removal_job.arr.tracker = AsyncMock() + removal_job.arr.tracker.extension_checked = {"checked-hash"} # Act processed_items = removal_job._get_items_to_process( # pylint: disable=W0212 diff --git a/tests/jobs/test_remove_unmonitored.py b/tests/jobs/test_remove_unmonitored.py index b633911..7bc3610 100644 --- a/tests/jobs/test_remove_unmonitored.py +++ b/tests/jobs/test_remove_unmonitored.py @@ -51,6 +51,16 @@ from src.jobs.remove_unmonitored import RemoveUnmonitored {101: False, 102: False}, ["1", "1"] ), + # One monitored, one not, one not matched yet + ( + [ + {"downloadId": "1", "detail_item_id": 101}, + {"downloadId": "2", "detail_item_id": 102}, + {"downloadId": "3", "detail_item_id": None} + ], + {101: True, 102: False}, + ["2"] + ), ] ) async def test_find_affected_items(queue_data, monitored_ids, expected_download_ids): diff --git a/tests/utils/test_queue_manager.py b/tests/utils/test_queue_manager.py index eb3b56e..e90b0d0 100644 --- a/tests/utils/test_queue_manager.py +++ b/tests/utils/test_queue_manager.py @@ -26,14 +26,16 @@ def test_format_queue_single_item(mock_queue_manager): "id": 1, } ] - expected = [{ - "downloadId": "abc123", - "downloadTitle": "Example Download Title", - "protocol": ["torrent"], - "status": ["queued"], - "IDs": [1], - }] - assert mock_queue_manager.format_queue(queue_items) == expected + expected = { + "abc123": { + "title": "Example Download Title", + "protocol": "torrent", + "status": "queued", + "queue_ids": [1], + } + } + result = mock_queue_manager.format_queue(queue_items) + assert result == expected def test_format_queue_multiple_same_download_id(mock_queue_manager): queue_items = [ @@ -52,14 +54,16 @@ def test_format_queue_multiple_same_download_id(mock_queue_manager): "id": 2, } ] - expected = [{ - "downloadId": "xyz789", - "downloadTitle": "Example Download Title", - "protocol": ["usenet"], - "status": ["downloading"], - "IDs": [1, 2], - }] - assert mock_queue_manager.format_queue(queue_items) == expected + expected = { + "xyz789": { + "title": "Example Download Title", + "protocol": "usenet", + "status": "downloading", + "queue_ids": [1, 2], + } + } + result = mock_queue_manager.format_queue(queue_items) + assert result == expected def test_format_queue_multiple_different_download_ids(mock_queue_manager): queue_items = [ @@ -78,20 +82,19 @@ def test_format_queue_multiple_different_download_ids(mock_queue_manager): "id": 20, } ] - expected = [ - { - "downloadId": "aaa111", - "downloadTitle": "Example Download Title A", - "protocol": ["torrent"], - "status": ["queued"], - "IDs": [10], + expected = { + 'aaa111': { + 'queue_ids': [10], + 'title': 'Example Download Title A', + 'protocol': 'torrent', + 'status': 'queued' }, - { - "downloadId": "bbb222", - "downloadTitle": "Example Download Title B", - "protocol": ["usenet"], - "status": ["completed"], - "IDs": [20], + 'bbb222': { + 'queue_ids': [20], + 'title': 'Example Download Title B', + 'protocol': 'usenet', + 'status': 'completed' } - ] - assert mock_queue_manager.format_queue(queue_items) == expected + } + result = mock_queue_manager.format_queue(queue_items) + assert result == expected