From 63d768bdc314dd90db203fe72a39b63e796bc80e Mon Sep 17 00:00:00 2001 From: Roland Date: Thu, 7 May 2026 17:03:12 -0600 Subject: [PATCH] Reporting download progress using download context. --- modules/videos/downloader.py | 13 +++++--- wrolpi/conftest.py | 9 ++++-- wrolpi/downloader.py | 44 +++++++++++++++++++++------- wrolpi/files/test/test_downloader.py | 14 +++++---- wrolpi/test/test_downloader.py | 40 +++++++++++++++---------- 5 files changed, 82 insertions(+), 38 deletions(-) diff --git a/modules/videos/downloader.py b/modules/videos/downloader.py index d20c7437..94731665 100755 --- a/modules/videos/downloader.py +++ b/modules/videos/downloader.py @@ -472,8 +472,11 @@ class ChannelDownloader(Downloader, ABC): during the file_worker await. """ _log_ytdlp_versions() + # Prefer ctx.extract_info when provided so tests can inject a fake without + # patching the module-level function. Falls back to the global for back-compat. + _extract = ctx.extract_info if ctx.extract_info else extract_info try: - info = extract_info(prepared.url, process=False) + info = _extract(prepared.url, process=False) except Exception as e: if _bot_blocked(e): raise BotBlockedDownloadError(str(e)) from e @@ -758,10 +761,12 @@ class VideoDownloader(Downloader, ABC): location = prepared.location # extract_info is sync but expensive; legacy held a session across it. We don't. + # Prefer ctx.extract_info for testability; fall back to the module global. + _extract = ctx.extract_info if ctx.extract_info else extract_info info_json = prepared.info_json if not info_json: try: - info_json = extract_info(url) + info_json = _extract(url) except yt_dlp.utils.DownloadError as e: if _bot_blocked(e): raise BotBlockedDownloadError(str(e)) from e @@ -863,7 +868,7 @@ class VideoDownloader(Downloader, ABC): if effective['sleep_requests'] > 0: cmd = (*cmd, '--sleep-requests', str(effective['sleep_requests'])) - on_stdout = make_progress_callback(download_for_log.id, parse_ytdlp_progress) + on_stdout = make_progress_callback(ctx.report_progress, parse_ytdlp_progress) try: if cookies_unlocked(): @@ -881,7 +886,7 @@ class VideoDownloader(Downloader, ABC): result = await self.process_runner(download_for_log, cmd, out_dir, debug=True, stdout_callback=on_stdout, ctx=ctx) finally: - clear_download_progress(download_for_log.id) + ctx.clear_progress() stdout = result.stdout.decode() stderr = result.stderr.decode() diff --git a/wrolpi/conftest.py b/wrolpi/conftest.py index dbab1f5f..33bcc7f1 100644 --- a/wrolpi/conftest.py +++ b/wrolpi/conftest.py @@ -380,12 +380,13 @@ def assert_downloads(test_session) -> Callable[[List[Dict]], None]: def make_test_ctx(*, is_cancelled=None, can_download=None, outside_download_window=None, run_command_=None, - download_timeout=None) -> DownloadContext: + download_timeout=None, + report_progress=None, clear_progress=None) -> DownloadContext: """Build a DownloadContext suitable for downloader unit tests. Defaults: never-cancelled, can_download True, no window restriction, real run_command, - no global timeout. Override any callable to simulate cancellation, throttling, or - fake subprocess output. + no global timeout, progress reporting is a no-op. Override any callable to simulate + cancellation, throttling, fake subprocess output, or to record progress events. Lives in wrolpi/conftest.py so any test (including primitive tests of process_runner / cancel_wrapper) can construct a ctx without crossing module boundaries. @@ -396,6 +397,8 @@ def make_test_ctx(*, is_cancelled=None, can_download=None, outside_download_window=outside_download_window or (lambda: False), run_command=run_command_ or run_command, download_timeout=download_timeout, + report_progress=report_progress or (lambda p: None), + clear_progress=clear_progress or (lambda: None), ) diff --git a/wrolpi/downloader.py b/wrolpi/downloader.py index 71ad3e25..446190f4 100644 --- a/wrolpi/downloader.py +++ b/wrolpi/downloader.py @@ -99,12 +99,13 @@ def clear_download_progress(download_id: int): api_app.shared_ctx.download_manager_data.update(data) -def make_progress_callback(download_id: int, parse_fn: callable) -> callable: - """Create a stdout callback that parses progress lines and updates shared_ctx. +def make_progress_callback(report: Callable[[dict], None], parse_fn: callable) -> callable: + """Create a stdout callback that parses progress lines and forwards them to `report`. The callback throttles updates to at most once per second. - :param download_id: The download ID to update progress for. + :param report: A function that takes a progress dict and persists it (production + wires this to set_download_progress; tests can pass a no-op or recorder). :param parse_fn: A function that takes a line string and returns a progress dict or None. """ last_update = [0.0] @@ -120,7 +121,7 @@ def make_progress_callback(download_id: int, parse_fn: callable) -> callable: progress = parse_fn(segment) if progress: last_update[0] = now_ - set_download_progress(download_id, progress) + report(progress) return return on_stdout @@ -176,16 +177,28 @@ class DownloadResult: class DownloadContext: """Per-download seam between DownloadManager and a Downloader. - Carries cancellation queries, the subprocess runner, and the effective timeout. - Production builds these as closures over api_app.shared_ctx and the singleton - DownloadManager (see DownloadContext.production); tests construct them directly - so a Downloader can run without Sanic's mp stack. + Carries cancellation queries, the subprocess runner, progress reporting, the effective + timeout, and the yt-dlp metadata fetcher. Production builds these as closures over + api_app.shared_ctx and the singleton DownloadManager (see DownloadContext.production); + tests construct them directly so a Downloader can run without Sanic's mp stack. """ is_cancelled: Callable[[], bool] can_download: Callable[[], bool] outside_download_window: Callable[[], bool] run_command: Callable[..., Awaitable[CommandResult]] download_timeout: Optional[int] + # Progress reporting: report writes a progress dict (parsed from subprocess stdout + # by make_progress_callback); clear removes it on completion or cancel. Production + # wires these to api_app.shared_ctx.download_manager_data['download_progress'] so + # the API server (potentially a different Sanic worker) can read it. Tests pass + # no-ops or recorders. + report_progress: Callable[[dict], None] + clear_progress: Callable[[], None] + # Optional: yt-dlp metadata fetcher. Only the videos module wires this; other + # downloaders ignore it. Tests can substitute a fake to bypass network I/O without + # patching the module-level extract_info function. None outside the videos path so + # the import isn't forced on every consumer. + extract_info: Optional[Callable[..., dict]] = None @classmethod def production(cls, manager: 'DownloadManager', download_id: int) -> 'DownloadContext': @@ -195,12 +208,23 @@ class DownloadContext: Cross-worker cancellation still flows through api_app.shared_ctx.download_manager_data['killed_downloads']; this just names the seam.""" + # Lazy import: extract_info lives in modules.videos and doesn't exist when the + # videos module is unavailable (e.g. trimmed deployments). We resolve it at + # context-construction time so the import is cheap and only happens once per + # download. + try: + from modules.videos.downloader import extract_info as _extract_info + except Exception: + _extract_info = None return cls( is_cancelled=lambda: download_id in api_app.shared_ctx.download_manager_data['killed_downloads'], can_download=lambda: manager.can_download, outside_download_window=lambda: manager.outside_download_window, run_command=run_command, download_timeout=get_wrolpi_config().download_timeout or None, + report_progress=lambda p: set_download_progress(download_id, p), + clear_progress=lambda: clear_download_progress(download_id), + extract_info=_extract_info, ) @@ -606,13 +630,13 @@ class Downloader: # the URL redirects to a CDN path that doesn't match the original filename. cmd = (*cmd, '-o', output_path.name) - on_stdout = make_progress_callback(download.id, parse_aria2c_progress) + on_stdout = make_progress_callback(ctx.report_progress, parse_aria2c_progress) try: result = await self.process_runner(download, cmd, destination, stdout_callback=on_stdout, ctx=ctx) finally: - clear_download_progress(download.id) + ctx.clear_progress() stderr = result.stderr.decode() if getattr(result, 'stderr', None) else '' stdout = result.stdout.decode() if getattr(result, 'stdout', None) else '' diff --git a/wrolpi/files/test/test_downloader.py b/wrolpi/files/test/test_downloader.py index 952f10e3..49fcfeee 100644 --- a/wrolpi/files/test/test_downloader.py +++ b/wrolpi/files/test/test_downloader.py @@ -144,12 +144,14 @@ async def test_file_downloader_aria2c_uses_output_flag(test_directory): fake_info = DownloadFileInfo(name='somefile.tar.gz', size=1024, status=200) + from wrolpi.conftest import make_test_ctx with mock.patch('wrolpi.downloader.get_download_info', mock.AsyncMock(return_value=fake_info)), \ - mock.patch.object(file_downloader, 'process_runner', fake_process_runner), \ - mock.patch('wrolpi.downloader.clear_download_progress'): + mock.patch.object(file_downloader, 'process_runner', fake_process_runner): download = mock.MagicMock() url = 'https://github.com/user/repo/releases/download/v1.0/somefile.tar.gz' - await file_downloader.download_file(download, url, destination) + # ctx.clear_progress (no-op in make_test_ctx) replaces the legacy + # mock.patch('wrolpi.downloader.clear_download_progress') hack. + await file_downloader.download_file(download, url, destination, ctx=make_test_ctx()) # The command should include -o with the expected filename. assert '-o' in captured_cmd @@ -181,13 +183,13 @@ async def test_file_downloader_meta4_no_output_flag(test_directory): fake_info = DownloadFileInfo(name='somefile.tar.gz', size=1024, status=200) + from wrolpi.conftest import make_test_ctx with mock.patch('wrolpi.downloader.get_download_info', mock.AsyncMock(return_value=fake_info)), \ mock.patch.object(file_downloader, 'get_meta4_contents', mock.AsyncMock(return_value=b'')), \ - mock.patch.object(file_downloader, 'process_runner', fake_process_runner), \ - mock.patch('wrolpi.downloader.clear_download_progress'): + mock.patch.object(file_downloader, 'process_runner', fake_process_runner): download = mock.MagicMock() url = 'https://example.com/somefile.tar.gz' - await file_downloader.download_file(download, url, destination) + await file_downloader.download_file(download, url, destination, ctx=make_test_ctx()) # When metalink is used, -o should NOT be in the command. assert '-o' not in captured_cmd diff --git a/wrolpi/test/test_downloader.py b/wrolpi/test/test_downloader.py index c98e5628..c4258daf 100644 --- a/wrolpi/test/test_downloader.py +++ b/wrolpi/test/test_downloader.py @@ -62,7 +62,13 @@ def test_parse_aria2c_progress_not_progress_line(): @pytest.mark.asyncio async def test_set_and_clear_download_progress(async_client): - """set_download_progress stores progress in shared_ctx and clear_download_progress removes it.""" + """set_download_progress stores progress in shared_ctx and clear_download_progress removes it. + + The two functions are the production wiring for ctx.report_progress / ctx.clear_progress; + this test asserts they actually move state through shared_ctx so the API server can + read it cross-worker. async_client is required because shared_ctx itself only exists + when the Sanic app is initialised. + """ from wrolpi.api_utils import api_app # No progress initially. data = dict(api_app.shared_ctx.download_manager_data) @@ -80,36 +86,40 @@ async def test_set_and_clear_download_progress(async_client): assert 42 not in data.get('download_progress', {}) -@pytest.mark.asyncio -async def test_make_progress_callback(async_client): - """make_progress_callback creates a callback that throttles and updates shared_ctx.""" - from wrolpi.api_utils import api_app +def test_make_progress_callback(): + """make_progress_callback throttles and forwards parsed progress to its `report` callable. + + Pure-logic test — no Sanic, no shared_ctx. We pass a list-append closure as `report` + and assert what landed in it. + """ + reported = [] def fake_parse(line): if 'progress' in line: return dict(percent=50) return None - callback = make_progress_callback(99, fake_parse) + callback = make_progress_callback(reported.append, fake_parse) # Non-matching line does nothing. callback('some other line') - data = dict(api_app.shared_ctx.download_manager_data) - assert 99 not in data.get('download_progress', {}) + assert reported == [] - # Matching line updates shared_ctx. + # Matching line invokes `report`. with mock.patch('wrolpi.downloader.time') as mock_time: mock_time.monotonic.return_value = 100.0 callback('progress line') - data = dict(api_app.shared_ctx.download_manager_data) - assert data['download_progress'][99] == dict(percent=50) + assert reported == [dict(percent=50)] - # A call within 1 second is throttled (uses a different progress value to verify). + # A call within 1 second is throttled. mock_time.monotonic.return_value = 100.5 - callback('progress line') # Would update, but throttled. + callback('progress line') + assert reported == [dict(percent=50)] # unchanged — throttle held - # Clean up. - clear_download_progress(99) + # After the throttle window passes, the next matching line is reported. + mock_time.monotonic.return_value = 102.0 + callback('progress line') + assert reported == [dict(percent=50), dict(percent=50)] @pytest.mark.asyncio