From 23df17b97e61fec8e4df4506a45cb257dfd33b93 Mon Sep 17 00:00:00 2001 From: Cody Logan Date: Wed, 1 Nov 2023 09:52:03 -0700 Subject: Refactor download processing code and add some tests --- src/wikiget/dl.py | 29 +++++++++++++++++++++++ src/wikiget/wikiget.py | 33 ++------------------------ tests/test_dl.py | 64 +++++++++++++++++++++++++++++++++++++++++++++++++- 3 files changed, 94 insertions(+), 32 deletions(-) diff --git a/src/wikiget/dl.py b/src/wikiget/dl.py index c0af6e9..20d8a07 100644 --- a/src/wikiget/dl.py +++ b/src/wikiget/dl.py @@ -49,6 +49,35 @@ def prep_download(dl: str, args: Namespace) -> File: return file +def process_download(args: Namespace) -> None: + if args.batch: + # batch download mode + errors = batch_download(args) + if errors: + # return non-zero exit code if any problems were encountered, even if some + # downloads completed successfully + logger.warning( + f"{errors} problem{'s'[:errors^1]} encountered during batch processing" + ) + sys.exit(1) # completed with errors + else: + # single download mode + try: + file = prep_download(args.FILE, args) + except ParseError as e: + logger.error(e) + sys.exit(1) + except FileExistsError as e: + logger.warning(e) + sys.exit(1) + except (ConnectionError, HTTPError, InvalidResponse, LoginError, APIError): + sys.exit(1) + + errors = download(file, args) + if errors: + sys.exit(1) # completed with errors + + def batch_download(args: Namespace) -> int: errors = 0 diff --git a/src/wikiget/wikiget.py b/src/wikiget/wikiget.py index f42da35..33f7b2c 100644 --- a/src/wikiget/wikiget.py +++ b/src/wikiget/wikiget.py @@ -17,14 +17,9 @@ import argparse import logging -import sys - -from mwclient import APIError, InvalidResponse, LoginError -from requests import ConnectionError, HTTPError import wikiget -from wikiget.dl import batch_download, download, prep_download -from wikiget.exceptions import ParseError +from wikiget.dl import process_download from wikiget.logging import configure_logging logger = logging.getLogger(__name__) @@ -134,28 +129,4 @@ def main() -> None: logger.info(f"Starting download session using wikiget {wikiget.__version__}") logger.debug(f"User agent: {wikiget.USER_AGENT}") - if args.batch: - # batch download mode - errors = batch_download(args) - if errors: - # return non-zero exit code if any problems were encountered, even if some - # downloads completed successfully - logger.warning( - f"{errors} problem{'s'[:errors^1]} encountered during batch processing" - ) - sys.exit(1) # completed with errors - else: - # single download mode - try: - file = prep_download(args.FILE, args) - except ParseError as e: - logger.error(e) - sys.exit(1) - except FileExistsError: - sys.exit(1) - except (ConnectionError, HTTPError, InvalidResponse, LoginError, APIError): - sys.exit(1) - - errors = download(file, args) - if errors: - sys.exit(1) # completed with errors + process_download(args) diff --git a/tests/test_dl.py b/tests/test_dl.py index 69ff2bb..61efbe0 100644 --- a/tests/test_dl.py +++ b/tests/test_dl.py @@ -16,10 +16,12 @@ # along with Wikiget. If not, see . from pathlib import Path +from unittest.mock import MagicMock, patch import pytest -from wikiget.dl import prep_download +from wikiget.dl import prep_download, process_download +from wikiget.file import File from wikiget.wikiget import construct_parser @@ -44,3 +46,63 @@ class TestPrepDownload: args = construct_parser().parse_args(["File:Example.jpg", "-o", str(tmp_file)]) with pytest.raises(FileExistsError): _ = prep_download(args.FILE, args) + + +class TestProcessDownload: + @patch("wikiget.dl.batch_download") + def test_batch_download(self, mock_batch_download: MagicMock) -> None: + """ + A successful batch download should not return any errors. + """ + args = construct_parser().parse_args(["-a", "batch.txt"]) + mock_batch_download.return_value = 0 + process_download(args) + assert mock_batch_download.called + + @patch("wikiget.dl.batch_download") + def test_batch_download_with_errors( + self, mock_batch_download: MagicMock, caplog: pytest.LogCaptureFixture + ) -> None: + """ + Any errors during batch download should create a log message containing the + number of errors and result in a non-zero exit code. + """ + args = construct_parser().parse_args(["-a", "batch.txt"]) + mock_batch_download.return_value = 4 + with pytest.raises(SystemExit) as e: + process_download(args) + assert mock_batch_download.called + assert e.value.code == 1 + assert "4 problems encountered during batch processing" in caplog.text + + @patch("wikiget.dl.prep_download") + @patch("wikiget.dl.download") + def test_single_download( + self, mock_download: MagicMock, mock_prep_download: MagicMock + ) -> None: + """ + A successful download should not return any errors. + """ + args = construct_parser().parse_args(["File:Example.jpg"]) + mock_download.return_value = 0 + mock_prep_download.return_value = MagicMock(File) + process_download(args) + assert mock_prep_download.called + assert mock_download.called + + @patch("wikiget.dl.prep_download") + @patch("wikiget.dl.download") + def test_single_download_with_errors( + self, mock_download: MagicMock, mock_prep_download: MagicMock + ) -> None: + """ + Any errors during download should result in a non-zero exit code. + """ + args = construct_parser().parse_args(["File:Example.jpg"]) + mock_download.return_value = 1 + mock_prep_download.return_value = MagicMock(File) + with pytest.raises(SystemExit) as e: + process_download(args) + assert mock_prep_download.called + assert mock_download.called + assert e.value.code == 1 -- cgit v1.2.3