mirror of
https://github.com/lrnselfreliance/wrolpi.git
synced 2026-10-26 10:51:36 -04:00
Reporting download progress using download context.
This commit is contained in:
parent
8f0bb78722
commit
63d768bdc3
5 changed files with 82 additions and 38 deletions
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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 ''
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue