From 8f608f2e647c185565021c15c0701d0276bd955d Mon Sep 17 00:00:00 2001 From: Atirna Date: Fri, 18 Sep 2026 01:21:20 +0530 Subject: [PATCH] fix(download): complete consumed Blackhole handoffs (#1345) A Blackhole watcher can consume the torrent before Shelfmark checks it, leaving the task in error even though the handoff succeeded. Complete the handoff when `add_download` successfully publishes the file, and stop requiring a `HandoffResult` path to remain present. Follow-up to #1312. ## Verification - A watcher that immediately reads and removes the torrent receives the exact bytes. The task changes from ERROR before this fix to COMPLETE afterward, without running book postprocessing. - The consumed-file regression fails on current main and passes here. Resident files, write failures, cancellation, magnet rejection and normal downloads remain covered: 81 focused tests pass. - Ruff lint and formatting pass for the changed files. --- shelfmark/download/clients/base_handler.py | 38 ++++++----- shelfmark/download/orchestrator.py | 8 --- tests/download/test_orchestrator_lifecycle.py | 65 +++++++++++++++---- 3 files changed, 76 insertions(+), 35 deletions(-) diff --git a/shelfmark/download/clients/base_handler.py b/shelfmark/download/clients/base_handler.py index e12dc41..f52ea20 100644 --- a/shelfmark/download/clients/base_handler.py +++ b/shelfmark/download/clients/base_handler.py @@ -871,6 +871,19 @@ class ExternalClientHandler(DownloadHandler, ABC): logger.info( "Added to %s: %s for '%s'", client.name, download_id, request.release_name ) + if getattr(client, "handoff_only", False) is True: + if cancel_flag.is_set(): + self._handle_cancelled_download( + client, download_id, request.protocol, status_callback + ) + return None + # A watcher can consume the publication as soon as add_download returns. + progress_callback(100) + self._on_download_complete(task) + return HandoffResult( + path=download_id, + message=f"Torrent file saved to {download_id}", + ) # Poll for progress return self._poll_and_complete( @@ -897,7 +910,7 @@ class ExternalClientHandler(DownloadHandler, ABC): cancel_flag: Event, progress_callback: Callable[[float], None], status_callback: Callable[[str, str | None], None], - ) -> str | HandoffResult | None: + ) -> str | None: """Poll the download client for progress and handle completion.""" poll_interval = self._poll_interval() # Track consecutive "not found" errors - torrents may take time to appear in client @@ -905,7 +918,7 @@ class ExternalClientHandler(DownloadHandler, ABC): max_not_found_retries = 15 # 15 retries * poll interval ~= 30s grace period try: - result: str | HandoffResult | None = None + result: str | None = None logger.debug("Starting poll for %s (content_type=%s)", download_id, task.content_type) while not cancel_flag.is_set(): status = client.get_status(download_id) @@ -1020,18 +1033,12 @@ class ExternalClientHandler(DownloadHandler, ABC): ) return None - if getattr(client, "handoff_only", False) is True: - result = HandoffResult( - path=str(source_path_obj), - message=f"Torrent file saved to {source_path_obj}", - ) - else: - result = self._handle_completed_file( - source_path=source_path_obj, - protocol=protocol, - task=task, - status_callback=status_callback, - ) + result = self._handle_completed_file( + source_path=source_path_obj, + protocol=protocol, + task=task, + status_callback=status_callback, + ) except Exception as e: logger.exception("Error during download polling") @@ -1042,8 +1049,7 @@ class ExternalClientHandler(DownloadHandler, ABC): # Clean up on success if result: self._on_download_complete(task) - if not isinstance(result, HandoffResult): - self._cleanup_refs[task.task_id] = (client, download_id, protocol) + self._cleanup_refs[task.task_id] = (client, download_id, protocol) return result diff --git a/shelfmark/download/orchestrator.py b/shelfmark/download/orchestrator.py index 6714156..845769e 100644 --- a/shelfmark/download/orchestrator.py +++ b/shelfmark/download/orchestrator.py @@ -756,14 +756,6 @@ def _download_task(task_id: str, cancel_flag: Event) -> str | None: if isinstance(temp_path, HandoffResult): handoff_path = Path(temp_path.path) - if not run_blocking_io(handoff_path.exists): - logger.error("Handler returned non-existent handoff path: %s", handoff_path) - _capture_task_error( - task, - message=f"Download file missing: {handoff_path}", - exc_type="MissingDownloadPath", - ) - return None status_callback("complete", temp_path.message) handler.post_process_cleanup(task, success=True) _clear_task_error_state(task) diff --git a/tests/download/test_orchestrator_lifecycle.py b/tests/download/test_orchestrator_lifecycle.py index 26b145a..011e823 100644 --- a/tests/download/test_orchestrator_lifecycle.py +++ b/tests/download/test_orchestrator_lifecycle.py @@ -1,12 +1,12 @@ from __future__ import annotations +from pathlib import Path from threading import Event from unittest.mock import ANY, MagicMock import pytest from shelfmark.core.models import DownloadTask -from shelfmark.release_sources import HandoffResult class _StopLoop(BaseException): @@ -231,27 +231,70 @@ def test_start_replaces_dead_coordinator_thread(monkeypatch): assert orchestrator._coordinator_thread is new_thread -def test_download_task_completes_blackhole_handoff_without_post_processing(monkeypatch, tmp_path): +@pytest.mark.parametrize("handoff", ["resident", "consumed", "write_error", "cancelled"]) +def test_download_task_completes_blackhole_handoff_without_post_processing( + monkeypatch, tmp_path, handoff +): import shelfmark.download.orchestrator as orchestrator + from shelfmark.download.clients import blackhole + from shelfmark.download.clients.base_handler import DownloadRequest + from shelfmark.download.clients.torrent_utils import TorrentInfo + from shelfmark.release_sources.prowlarr.handler import ProwlarrHandler - handoff_file = tmp_path / "release.torrent" - handoff_file.write_bytes(b"torrent-bytes") + handoff_file = tmp_path / "Book.torrent" + cancel = Event() task = DownloadTask(task_id="blackhole-task", source="prowlarr", title="Book") queue = MagicMock() queue.get_task.return_value = task - handler = MagicMock() - handler.download.return_value = HandoffResult( - path=str(handoff_file), - message=f"Torrent file saved to {handoff_file}", + monkeypatch.setattr( + blackhole.config, + "get", + lambda key, default=None: str(tmp_path) if key == "BLACKHOLE_DIRECTORY" else default, ) + monkeypatch.setattr( + blackhole, + "extract_torrent_info", + lambda *_args, **_kwargs: TorrentInfo("abc123", b"torrent-bytes", False), + ) + client = blackhole.BlackholeClient() + handler = ProwlarrHandler() + monkeypatch.setattr(handler, "_get_client", lambda _protocol: client) + monkeypatch.setattr( + handler, + "_resolve_download", + lambda *_args: DownloadRequest( + "https://indexer.example/book.torrent", "torrent", "Book", None + ), + ) + replace = Path.replace + consumed = [] + def publish(path, target): + if handoff == "write_error": + raise OSError("handoff directory is not writable") + result = replace(path, target) + if handoff == "consumed": + consumed.append(target.read_bytes()) + target.unlink() + elif handoff == "cancelled": + cancel.set() + return result + + monkeypatch.setattr(Path, "replace", publish) monkeypatch.setattr(orchestrator, "book_queue", queue) monkeypatch.setattr(orchestrator, "get_handler", lambda _source: handler) monkeypatch.setattr(orchestrator, "_source_unavailable_message", lambda _source: None) monkeypatch.setattr(orchestrator, "post_process_download", MagicMock()) - result = orchestrator._download_task(task.task_id, Event()) + result = orchestrator._download_task(task.task_id, cancel) - assert result == str(handoff_file) + assert result == (None if handoff in ("write_error", "cancelled") else str(handoff_file)) + assert handoff_file.exists() is (handoff in ("resident", "cancelled")) + assert consumed == ([b"torrent-bytes"] if handoff == "consumed" else []) + assert (task.last_error_message is not None) is (handoff == "write_error") + assert not list(tmp_path.glob(".blackhole-*")) orchestrator.post_process_download.assert_not_called() - handler.post_process_cleanup.assert_called_once_with(task, success=True) + if result: + queue.update_progress.assert_called_once_with(task.task_id, 100) + else: + queue.update_progress.assert_not_called()