From f2dc249e01099414720eb9790aa21deba6b5eeda Mon Sep 17 00:00:00 2001 From: Roland Date: Wed, 6 May 2026 19:39:07 -0600 Subject: [PATCH] Migrating Kiwix downloaders. --- modules/zim/downloader.py | 126 ++++++++++++++--- modules/zim/test/test_zim_downloaders.py | 166 +++++++++++++++++++++++ wrolpi/conftest.py | 3 +- 3 files changed, 272 insertions(+), 23 deletions(-) create mode 100644 modules/zim/test/test_zim_downloaders.py diff --git a/modules/zim/downloader.py b/modules/zim/downloader.py index ff515a04..8d2fd8df 100644 --- a/modules/zim/downloader.py +++ b/modules/zim/downloader.py @@ -1,11 +1,20 @@ +import pathlib +from dataclasses import dataclass from http import HTTPStatus -from typing import List +from typing import List, Optional + +from sqlalchemy.orm import Session from modules.zim import lib from wrolpi.common import logger, aiohttp_get, get_html_soup -from wrolpi.downloader import Downloader, Download, DownloadResult +from wrolpi.downloader import Downloader, Download, DownloadContext, DownloadResult -__all__ = ['KiwixCatalogDownloader', 'KiwixZimDownloader', 'kiwix_zim_downloader', 'kiwix_catalog_downloader'] +__all__ = [ + 'KiwixCatalogDownloader', 'KiwixZimDownloader', + 'PreparedKiwixCatalog', 'ExecutedKiwixCatalog', + 'PreparedKiwixZim', 'ExecutedKiwixZim', + 'kiwix_zim_downloader', 'kiwix_catalog_downloader', +] from wrolpi.files.lib import upsert_file from wrolpi.vars import DOWNLOAD_USER_AGENT @@ -48,6 +57,33 @@ async def fetch_hrefs(url: str) -> List[str]: return downloads +# --------------------------------------------------------------------------- +# KiwixCatalogDownloader +# --------------------------------------------------------------------------- + + +@dataclass +class PreparedKiwixCatalog: + """Plan produced by KiwixCatalogDownloader.prepare_download. + + Splits the URL into the parent listing URL and the basename prefix used to match + Zim files in that listing. + """ + url: str + parent_url: str + name: str + + +@dataclass +class ExecutedKiwixCatalog: + """Output of KiwixCatalogDownloader.execute_download. + + `latest_zim_url` is None when the listing contained no matching files; finalize + converts that into a failure DownloadResult. + """ + latest_zim_url: Optional[str] + + class KiwixCatalogDownloader(Downloader): """ Downloads the catalog of Zim files from Kiwix. Searches for the latest Zim that matches the provided URL. @@ -58,59 +94,105 @@ class KiwixCatalogDownloader(Downloader): name = 'kiwix_catalog' listable = False - async def do_download(self, download: Download) -> DownloadResult: + def prepare_download(self, session: Session, download: Download) -> PreparedKiwixCatalog: + """Split the download URL into the parent listing URL and the basename prefix.""" url = download.url - - # Get the parent directory of the file to be downloaded *parents, name = url.split('/') parent_url = '/'.join(parents) + return PreparedKiwixCatalog(url=url, parent_url=parent_url, name=name) - # Find the latest Zim that has the correct name. - downloadable_files = await fetch_hrefs(parent_url) - matching_zims = sorted([i for i in downloadable_files if i.startswith(name) and i.endswith('.zim')]) + async def execute_download(self, prepared: PreparedKiwixCatalog, ctx: DownloadContext, + download: Download = None) -> ExecutedKiwixCatalog: + """Fetch the parent directory listing and pick the latest matching Zim.""" + downloadable_files = await fetch_hrefs(prepared.parent_url) + matching_zims = sorted([i for i in downloadable_files + if i.startswith(prepared.name) and i.endswith('.zim')]) if not matching_zims: - return DownloadResult(success=False, error=f'No Zim files match {str(repr(url))}') + return ExecutedKiwixCatalog(latest_zim_url=None) logger.debug(f'Found matching Zims: {matching_zims}') - latest_zim = f'{parent_url}/{matching_zims[-1]}' + return ExecutedKiwixCatalog(latest_zim_url=f'{prepared.parent_url}/{matching_zims[-1]}') - return DownloadResult( - success=True, - downloads=[latest_zim, ], - ) + def finalize_download(self, session: Session, download: Download, + executed: ExecutedKiwixCatalog) -> DownloadResult: + if executed.latest_zim_url is None: + return DownloadResult(success=False, error=f'No Zim files match {str(repr(download.url))}') + return DownloadResult(success=True, downloads=[executed.latest_zim_url]) kiwix_catalog_downloader = KiwixCatalogDownloader() +# --------------------------------------------------------------------------- +# KiwixZimDownloader +# --------------------------------------------------------------------------- + + +@dataclass +class PreparedKiwixZim: + """Plan produced by KiwixZimDownloader.prepare_download: zim_directory exists on disk.""" + url: str + zim_directory: pathlib.Path + + +@dataclass +class ExecutedKiwixZim: + """Output of KiwixZimDownloader.execute_download. + + `is_valid` is False when zimcheck flagged the file as corrupt; finalize turns that + into a failure DownloadResult. When zimcheck is not installed (FileNotFoundError), + we treat the file as valid (matching legacy behaviour). + """ + output_path: pathlib.Path + is_valid: bool + + class KiwixZimDownloader(Downloader): """Downloads a Zim file to the zim directory.""" name = 'kiwix_zim' listable = False - async def do_download(self, download: Download) -> DownloadResult: + def prepare_download(self, session: Session, download: Download) -> PreparedKiwixZim: zim_directory = lib.get_zim_directory() zim_directory.mkdir(parents=True, exist_ok=True) + return PreparedKiwixZim(url=download.url, zim_directory=zim_directory) - output_path = await self.download_file(download, download.url, zim_directory) + async def execute_download(self, prepared: PreparedKiwixZim, ctx: DownloadContext, + download: Download = None) -> ExecutedKiwixZim: + """Download the Zim file, flag any outdated zims, validate via zimcheck, and + register the new file via upsert_file when valid. + + Note: lib.flag_outdated_zim_files (sets a Sanic flag) and upsert_file (opens its + own DB session) are pre-existing side effects we keep on the post-download path. + Splitting upsert_file is a future PR shared with the FileDownloader migration. + """ + download = download if download is not None else Download(url=prepared.url) + output_path = await self.download_file(download, prepared.url, prepared.zim_directory, ctx=ctx) # Notify the maintainer if outdated Zim files are lying around. lib.flag_outdated_zim_files() + is_valid = True try: return_code = await lib.check_zim(output_path) if return_code > 0: # zimcheck ran, but the file is invalid. - return DownloadResult(success=False, error=f'Zim file is invalid') + is_valid = False except FileNotFoundError: logger.warning(f'Not validating {output_path} because zimcheck is not installed') - # Add the new Zim to the FileGroups, model it. - # The model_zim() function triggers a Kiwix restart via switch when a new Zim is created. - await upsert_file(output_path) + if is_valid: + # Add the new Zim to the FileGroups, model it. + # The model_zim() function triggers a Kiwix restart via switch when a new Zim is created. + await upsert_file(output_path) + logger.info(f'Successfully downloaded Zim {prepared.url} to {output_path}') - logger.info(f'Successfully downloaded Zim {download.url} to {output_path}') + return ExecutedKiwixZim(output_path=output_path, is_valid=is_valid) + def finalize_download(self, session: Session, download: Download, + executed: ExecutedKiwixZim) -> DownloadResult: + if not executed.is_valid: + return DownloadResult(success=False, error='Zim file is invalid') # Location is just the generic Kiwix viewer. return DownloadResult(success=True, location='/zim') diff --git a/modules/zim/test/test_zim_downloaders.py b/modules/zim/test/test_zim_downloaders.py new file mode 100644 index 00000000..5ae33b4e --- /dev/null +++ b/modules/zim/test/test_zim_downloaders.py @@ -0,0 +1,166 @@ +"""Phase-split tests for KiwixCatalogDownloader and KiwixZimDownloader. + +These exercise prepare_download / finalize_download (and execute_download for the +catalog, where it's pure HTTP) in isolation — no async_client, no test_download_manager, +no real subprocess. The Zim downloader's execute phase calls into download_file, +flag_outdated_zim_files, check_zim (subprocess), and upsert_file (own DB session); a +unit test of that full pipeline would effectively be an integration test, so it's +deliberately omitted. +""" +from unittest import mock + +import pytest + +from modules.archive.conftest import make_test_ctx +from modules.zim.downloader import ( + ExecutedKiwixCatalog, + ExecutedKiwixZim, + PreparedKiwixCatalog, + PreparedKiwixZim, + kiwix_catalog_downloader, + kiwix_zim_downloader, +) +from wrolpi.downloader import Download + + +# --------------------------------------------------------------------------- +# KiwixCatalogDownloader +# --------------------------------------------------------------------------- + + +def test_catalog_prepare_extracts_parent_and_name(test_session): + """prepare_download splits the URL into the listing parent and the filename prefix.""" + download = Download( + url='https://download.kiwix.org/zim/wikipedia/wikipedia_en_all_nopic', + downloader='kiwix_catalog', + ) + + prepared = kiwix_catalog_downloader.prepare_download(test_session, download) + + assert isinstance(prepared, PreparedKiwixCatalog) + assert prepared.url == download.url + assert prepared.parent_url == 'https://download.kiwix.org/zim/wikipedia' + assert prepared.name == 'wikipedia_en_all_nopic' + + +@pytest.mark.asyncio +async def test_catalog_execute_picks_latest_matching_zim(test_session): + """execute_download fetches the parent listing and picks the lexicographically last + matching .zim — sorted ascending, so the newest YYYY-MM Kiwix filename wins.""" + listing = [ + 'wikipedia_en_all_nopic_2024-01.zim', + 'wikipedia_en_all_nopic_2024-06.zim', + 'wikipedia_en_all_nopic_2024-12.zim', + 'wikipedia_en_all_nopic_2024-12.zim.meta4', # not a .zim file → skipped + 'unrelated_corpus_2024-12.zim', # different prefix → skipped + ] + + async def fake_fetch_hrefs(url): + assert url == 'https://download.kiwix.org/zim/wikipedia' + return listing + + prepared = PreparedKiwixCatalog( + url='https://download.kiwix.org/zim/wikipedia/wikipedia_en_all_nopic', + parent_url='https://download.kiwix.org/zim/wikipedia', + name='wikipedia_en_all_nopic', + ) + + with mock.patch('modules.zim.downloader.fetch_hrefs', fake_fetch_hrefs): + executed = await kiwix_catalog_downloader.execute_download(prepared, make_test_ctx()) + + assert isinstance(executed, ExecutedKiwixCatalog) + assert executed.latest_zim_url == \ + 'https://download.kiwix.org/zim/wikipedia/wikipedia_en_all_nopic_2024-12.zim' + + +@pytest.mark.asyncio +async def test_catalog_execute_no_matches(test_session): + """When the listing has no matching files, executed.latest_zim_url is None.""" + + async def fake_fetch_hrefs(url): + return ['some_other_corpus_2024-01.zim'] + + prepared = PreparedKiwixCatalog( + url='https://download.kiwix.org/zim/wikipedia/wikipedia_en_all_nopic', + parent_url='https://download.kiwix.org/zim/wikipedia', + name='wikipedia_en_all_nopic', + ) + + with mock.patch('modules.zim.downloader.fetch_hrefs', fake_fetch_hrefs): + executed = await kiwix_catalog_downloader.execute_download(prepared, make_test_ctx()) + + assert executed.latest_zim_url is None + + +def test_catalog_finalize_no_match_returns_failure(test_session): + """latest_zim_url=None → DownloadResult(success=False) with the URL in the error.""" + download = Download(url='https://example.com/zim/foo', downloader='kiwix_catalog') + executed = ExecutedKiwixCatalog(latest_zim_url=None) + + result = kiwix_catalog_downloader.finalize_download(test_session, download, executed) + + assert result.success is False + assert 'https://example.com/zim/foo' in (result.error or '') + + +def test_catalog_finalize_with_match_emits_child_download(test_session): + """The latest matching URL is emitted as a child download for the kiwix_zim downloader.""" + download = Download(url='https://example.com/zim/foo', downloader='kiwix_catalog') + executed = ExecutedKiwixCatalog( + latest_zim_url='https://example.com/zim/foo_2024-12.zim', + ) + + result = kiwix_catalog_downloader.finalize_download(test_session, download, executed) + + assert result.success is True + assert result.downloads == ['https://example.com/zim/foo_2024-12.zim'] + + +# --------------------------------------------------------------------------- +# KiwixZimDownloader +# --------------------------------------------------------------------------- + + +def test_zim_prepare_creates_zim_directory(test_session, test_directory, monkeypatch): + """prepare_download ensures the zim directory exists and returns it on the plan.""" + expected = test_directory / 'zims' + assert not expected.exists() + + # Override get_zim_directory so we don't depend on wrolpi config defaults. + monkeypatch.setattr('modules.zim.lib.get_zim_directory', lambda: expected) + + download = Download( + url='https://example.com/zim/wikipedia_en_all_nopic_2024-12.zim', + downloader='kiwix_zim', + ) + + prepared = kiwix_zim_downloader.prepare_download(test_session, download) + + assert isinstance(prepared, PreparedKiwixZim) + assert prepared.url == download.url + assert prepared.zim_directory == expected + assert expected.is_dir() + + +def test_zim_finalize_invalid_zim_returns_failure(test_session, test_directory): + """is_valid=False → DownloadResult(success=False) with the legacy error message.""" + output_path = test_directory / 'zims/wikipedia_en_all_nopic_2024-12.zim' + executed = ExecutedKiwixZim(output_path=output_path, is_valid=False) + download = Download(url='https://example.com/zim/foo', downloader='kiwix_zim') + + result = kiwix_zim_downloader.finalize_download(test_session, download, executed) + + assert result.success is False + assert result.error == 'Zim file is invalid' + + +def test_zim_finalize_valid_returns_success(test_session, test_directory): + """is_valid=True → DownloadResult(success=True) with the generic Kiwix viewer location.""" + output_path = test_directory / 'zims/wikipedia_en_all_nopic_2024-12.zim' + executed = ExecutedKiwixZim(output_path=output_path, is_valid=True) + download = Download(url='https://example.com/zim/foo', downloader='kiwix_zim') + + result = kiwix_zim_downloader.finalize_download(test_session, download, executed) + + assert result.success is True + assert result.location == '/zim' diff --git a/wrolpi/conftest.py b/wrolpi/conftest.py index 7ef26760..8e292872 100644 --- a/wrolpi/conftest.py +++ b/wrolpi/conftest.py @@ -1059,7 +1059,8 @@ def mock_downloader_download_file(): """Mock the Downloader.download_file method.""" contents = None - async def _download_file(self, id_, url, destination): + async def _download_file(self, id_, url, destination, **kwargs): + # **kwargs absorbs ctx= and any future optional args added to Downloader.download_file. name = url.split('/')[-1] output_path = destination / name output_path.write_bytes(contents or b'')