Reporting download progress using download context.

This commit is contained in:
Roland 2026-05-07 17:03:12 -06:00 • committed by lrnselfreliance
parent 8f0bb78722
commit 63d768bdc3
5 changed files with 82 additions and 38 deletions

View file

@ -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()

View file

@ -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),
)

View file

@ -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 ''

View file

@ -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'<metalink/>')), \
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

View file

@ -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