diff --git a/main.py b/main.py
index 3be0cf2c..2d4c420c 100755
--- a/main.py
+++ b/main.py
@@ -7,12 +7,13 @@ import sys
import pytz
from sanic import Sanic
+from sanic.signals import Event
from wrolpi import root_api, BEFORE_STARTUP_FUNCTIONS, after_startup, limit_concurrent, admin
from wrolpi.common import logger, get_config, import_modules, check_media_directory
from wrolpi.dates import set_timezone
from wrolpi.downloader import download_manager
-from wrolpi.vars import PROJECT_DIR, DOCKERIZED
+from wrolpi.vars import PROJECT_DIR, DOCKERIZED, PYTEST
from wrolpi.version import get_version_string
logger = logger.getChild('wrolpi-main')
@@ -191,16 +192,21 @@ def periodic_downloads(app: Sanic, loop):
logger.info('Starting download manager.')
async def _periodic_download():
- download_manager.renew_recurring_downloads()
- download_manager.delete_old_once_downloads()
+ download_manager.start_workers()
await download_manager.do_downloads()
- seconds = 60
- await asyncio.sleep(seconds)
- app.add_task(_periodic_download())
+ await asyncio.sleep(60)
+ app.add_task(_periodic_download()) # noqa
app.add_task(_periodic_download())
+@root_api.api_app.signal(Event.SERVER_SHUTDOWN_BEFORE)
+def handle_server_shutdown(*args, **kwargs):
+ """Stop downloads when server is shutting down."""
+ if not PYTEST:
+ download_manager.stop()
+
+
if __name__ == '__main__':
loop_ = asyncio.get_event_loop()
result_ = loop_.run_until_complete(main(loop_))
diff --git a/modules/archive/__init__.py b/modules/archive/__init__.py
index 1a0afe57..1d059e8a 100644
--- a/modules/archive/__init__.py
+++ b/modules/archive/__init__.py
@@ -30,17 +30,16 @@ class ArchiveDownloader(Downloader, ABC):
"""
return True, None
- def do_download(self, download: Download) -> DownloadResult:
+ async def do_download(self, download: Download) -> DownloadResult:
if download.attempts > 3:
raise UnrecoverableDownloadError(f'Max download attempts reached for {download.url}')
- archive = lib.do_archive(download.url)
+ archive = await lib.do_archive(download.url)
return DownloadResult(success=True, location=f'/archive/{archive.id}')
@optional_session
- def already_downloaded(self, url: str, session: Session = None) -> Optional[Archive]:
- archive = session.query(Archive).filter_by(url=url).one_or_none()
- return archive
+ def already_downloaded(self, url: str, session: Session = None) -> bool:
+ return bool(session.query(Archive).filter_by(url=url).count())
# Archive downloader is the last downloader which should be used.
diff --git a/modules/archive/api.py b/modules/archive/api.py
index b8a62e28..de25c386 100644
--- a/modules/archive/api.py
+++ b/modules/archive/api.py
@@ -1,8 +1,7 @@
import asyncio
from http import HTTPStatus
-from requests import Request
-from sanic import response
+from sanic import response, Request
from sanic_ext import validate
from sanic_ext.extensions.openapi import openapi
diff --git a/modules/archive/lib.py b/modules/archive/lib.py
index 7c9590c2..1a325a96 100644
--- a/modules/archive/lib.py
+++ b/modules/archive/lib.py
@@ -10,14 +10,15 @@ from datetime import datetime
from itertools import groupby
from typing import Iterator, Optional, Tuple, List
-import requests
+import aiohttp
from bs4 import BeautifulSoup
from selenium import webdriver
from sqlalchemy.orm import Session
from modules.archive.models import Domain, Archive
from wrolpi.cmd import which
-from wrolpi.common import get_media_directory, logger, chunks, extract_domain, chdir, escape_file_name, walk
+from wrolpi.common import get_media_directory, logger, chunks, extract_domain, chdir, escape_file_name, walk, \
+ aiohttp_post
from wrolpi.dates import now, Seconds, local_timezone
from wrolpi.db import get_db_session, get_db_curs, get_ranked_models, optional_session
from wrolpi.errors import InvalidDomain, UnknownURL, InvalidArchive
@@ -102,7 +103,7 @@ def get_new_archive_files(url: str, title: Optional[str]) -> ArchiveFiles:
ARCHIVE_TIMEOUT = Seconds.minute * 10 # Wait at most 10 minutes for response.
-def request_archive(url: str):
+async def request_archive(url: str):
"""
Send a request to the archive service to archive the URL.
"""
@@ -110,18 +111,17 @@ def request_archive(url: str):
data = {'url': url}
try:
- resp = requests.post(f'{ARCHIVE_SERVICE}/json', json=data, timeout=ARCHIVE_TIMEOUT)
+ contents, status = await aiohttp_post(f'{ARCHIVE_SERVICE}/json', json_=data, timeout=ARCHIVE_TIMEOUT)
+ readability = contents['readability']
+ # Compressed base64
+ singlefile = contents['singlefile']
+ screenshot = contents['screenshot']
+
+ logger.debug(f'archive request status code {status}')
except Exception as e:
logger.error('Error when requesting archive', exc_info=e)
raise
- logger.debug(f'archive request status {resp.status_code=}')
-
- readability = resp.json()['readability']
- # Compressed base64
- singlefile = resp.json()['singlefile']
- screenshot = resp.json()['screenshot']
-
if not (screenshot or singlefile or readability):
raise Exception('singlefile response was empty!')
@@ -147,8 +147,8 @@ def request_archive(url: str):
SINGLE_FILE_BIN = which('single-file',
- '/usr/bin/single-file', # rpi ubuntu
- '/usr/local/bin/single-file', # debian
+ '/usr/bin/single-file', # rpi ubuntu
+ '/usr/local/bin/single-file', # debian
warn=True)
CHROMIUM = which('chromium-browser', 'chromium',
'/usr/bin/chromium-browser', # rpi ubuntu
@@ -194,8 +194,8 @@ def local_screenshot(url: str) -> bytes:
READABILITY_BIN = which('readability-extractor',
- '/usr/bin/readability-extractor', # rpi ubuntu
- '/usr/local/bin/readability-extractor', # debian
+ '/usr/bin/readability-extractor', # rpi ubuntu
+ '/usr/local/bin/readability-extractor', # debian
warn=True)
@@ -227,7 +227,7 @@ def local_archive(url: str):
return singlefile, readability, screenshot
-def do_archive(url: str) -> Archive:
+async def do_archive(url: str) -> Archive:
"""
Perform the real archive request to the archiving service. Store the resulting data into files. Create an Archive
record in the DB. Create Domain/URL if missing.
@@ -236,7 +236,7 @@ def do_archive(url: str) -> Archive:
if DOCKERIZED or PYTEST:
# Perform the archive in the Archive docker container. (Typically in the development environment).
- singlefile, readability, screenshot = request_archive(url)
+ singlefile, readability, screenshot = await request_archive(url)
else:
# Perform the archive using locally installed executables.
singlefile, readability, screenshot = local_archive(url)
diff --git a/modules/archive/test/test_lib.py b/modules/archive/test/test_lib.py
index 2bd134dd..560272fa 100644
--- a/modules/archive/test/test_lib.py
+++ b/modules/archive/test/test_lib.py
@@ -14,11 +14,11 @@ from wrolpi.dates import local_timezone
from wrolpi.db import get_db_session
from wrolpi.media_path import MediaPath
from wrolpi.root_api import CustomJSONEncoder
-from wrolpi.test.common import TestAPI, wrap_test_db, skip_circleci
+from wrolpi.test.common import skip_circleci
def make_fake_request_archive(readability=True, screenshot=True, title=True):
- def fake_request_archive(_):
+ async def fake_request_archive(_):
singlefile = '\ntest single-file\nジにてこちら\nsome title'
r = dict(
content=f'test readability content',
@@ -31,63 +31,61 @@ def make_fake_request_archive(readability=True, screenshot=True, title=True):
return fake_request_archive
-class TestArchive(TestAPI):
+@pytest.mark.asyncio
+async def test_no_screenshot(test_session):
+ with mock.patch('modules.archive.lib.request_archive', make_fake_request_archive(screenshot=False)):
+ archive = await do_archive('https://example.com')
+ assert isinstance(archive.singlefile_path, MediaPath)
+ assert isinstance(archive.readability_path, MediaPath)
+ # Screenshot was empty
+ assert archive.screenshot_path is None
- def setUp(self) -> None:
- super().setUp()
- (pathlib.Path(self.tmp_dir.name) / 'archive').mkdir(exist_ok=True)
- @wrap_test_db
- def test_no_screenshot(self):
- with mock.patch('modules.archive.lib.request_archive', make_fake_request_archive(screenshot=False)):
- archive = do_archive('https://example.com')
- self.assertIsInstance(archive.singlefile_path, MediaPath)
- self.assertIsInstance(archive.readability_path, MediaPath)
- # Screenshot was empty
- self.assertIsNone(archive.screenshot_path)
+@pytest.mark.asyncio
+async def test_no_readability(test_session):
+ with mock.patch('modules.archive.lib.request_archive', make_fake_request_archive(readability=False)):
+ archive = await do_archive('https://example.com')
+ assert isinstance(archive.singlefile_path, MediaPath)
+ assert isinstance(archive.screenshot_path, MediaPath)
+ # Readability empty
+ assert archive.readability_path is None
+ assert archive.readability_txt_path is None
- @wrap_test_db
- def test_no_readability(self):
- with mock.patch('modules.archive.lib.request_archive', make_fake_request_archive(readability=False)):
- archive = do_archive('https://example.com')
- self.assertIsInstance(archive.singlefile_path, MediaPath)
- self.assertIsInstance(archive.screenshot_path, MediaPath)
- # Readability empty
- self.assertIsNone(archive.readability_path)
- self.assertIsNone(archive.readability_txt_path)
- @wrap_test_db
- def test_dict(self):
- with mock.patch('modules.archive.lib.request_archive', make_fake_request_archive()):
- d = do_archive('https://example.com').dict()
- self.assertIsInstance(d, dict)
- json.dumps(d, cls=CustomJSONEncoder)
+@pytest.mark.asyncio
+async def test_dict(test_session):
+ with mock.patch('modules.archive.lib.request_archive', make_fake_request_archive()):
+ d = (await do_archive('https://example.com')).dict()
+ assert isinstance(d, dict)
+ json.dumps(d, cls=CustomJSONEncoder)
- @wrap_test_db
- def test_relationships(self):
- with get_db_session(commit=True) as session:
- url = 'https://wrolpi.org:443'
- domain = get_or_create_domain(session, url)
- archive = Archive(
- singlefile_path=f'{self.tmp_dir.name}/wrolpi.org:443/foo',
- title='bar',
- url=url,
- domain_id=domain.id,
- )
- session.add(archive)
- session.flush()
- self.assertEqual(archive.domain, domain)
+@pytest.mark.asyncio
+async def test_relationships(test_session, test_directory):
+ with get_db_session(commit=True) as session:
+ url = 'https://wrolpi.org:443'
+ domain = get_or_create_domain(session, url)
+ archive = Archive(
+ singlefile_path=f'{test_directory.name}/wrolpi.org:443/foo',
+ title='bar',
+ url=url,
+ domain_id=domain.id,
+ )
+ session.add(archive)
+ session.flush()
- @wrap_test_db
- def test_validate_paths(self):
- with mock.patch('modules.archive.lib.request_archive', make_fake_request_archive()):
- archive = do_archive('https://example.com')
- try:
- with get_db_session(commit=True):
- archive.singlefile_path = 'asdf'
- except ValueError as e:
- self.assertIn('relative', str(e), f'Relative path error was not raised')
+ assert archive.domain == domain
+
+
+@pytest.mark.asyncio
+async def test_validate_paths(test_session):
+ with mock.patch('modules.archive.lib.request_archive', make_fake_request_archive()):
+ archive = await do_archive('https://example.com')
+ try:
+ with get_db_session(commit=True):
+ archive.singlefile_path = 'asdf'
+ except ValueError as e:
+ assert 'relative' in str(e), f'Relative path error was not raised'
def test_archive_refresh_deleted_archive(test_session, archive_directory, archive_factory):
@@ -225,10 +223,11 @@ def test_get_domains(test_session, archive_factory):
assert [i['domain'] for i in get_domains()] == []
-def test_new_archive(test_session, fake_now):
+@pytest.mark.asycio
+async def test_new_archive(test_session, fake_now):
with mock.patch('modules.archive.lib.request_archive', make_fake_request_archive()):
fake_now(datetime(2000, 1, 1))
- archive1 = do_archive('https://example.com')
+ archive1 = await do_archive('https://example.com')
# Everything is filled out.
assert isinstance(archive1, Archive)
assert archive1.archive_datetime is not None
@@ -251,18 +250,19 @@ def test_new_archive(test_session, fake_now):
assert json.load(fh) == {'title': 'ジにてこちら', 'url': 'https://example.com'}
fake_now(datetime(2000, 1, 2))
- archive2 = do_archive('https://example.com')
+ archive2 = await do_archive('https://example.com')
# Domain is reused.
assert archive1.domain == archive2.domain
-def test_get_title_from_html(test_session, fake_now):
+@pytest.mark.asyncio
+async def test_get_title_from_html(test_session, fake_now):
fake_now(datetime(2000, 1, 1))
with mock.patch('modules.archive.lib.request_archive', make_fake_request_archive()):
- archive = do_archive('example.com')
+ archive = await do_archive('example.com')
assert archive.title == 'ジにてこちら'
- def fake_request_archive(_):
+ async def fake_request_archive(_):
singlefile = '\ntest single-file\nジにてこちら\nsome title'
r = dict(
content=f'test readability content',
@@ -273,10 +273,10 @@ def test_get_title_from_html(test_session, fake_now):
fake_now(datetime(2000, 1, 2))
with mock.patch('modules.archive.lib.request_archive', fake_request_archive):
- archive = do_archive('example.com')
+ archive = await do_archive('example.com')
assert archive.title == 'some title'
- def fake_request_archive(_):
+ async def fake_request_archive(_):
singlefile = ''
r = dict(
content=f'missing a title',
@@ -287,7 +287,7 @@ def test_get_title_from_html(test_session, fake_now):
fake_now(datetime(2000, 1, 3))
with mock.patch('modules.archive.lib.request_archive', fake_request_archive):
- archive = do_archive('example.com')
+ archive = await do_archive('example.com')
assert archive.title is None
@@ -313,14 +313,15 @@ def test_get_new_archive_files(fake_now):
@skip_circleci
-def test_title_in_filename(test_session, fake_now, test_directory):
+@pytest.mark.asyncio
+async def test_title_in_filename(test_session, fake_now, test_directory):
"""
The Archive files have the title in the path.
"""
fake_now(datetime(2000, 1, 1))
with mock.patch('modules.archive.lib.request_archive', make_fake_request_archive()):
- archive1 = do_archive('example.com')
+ archive1 = await do_archive('example.com')
assert archive1.title == 'ジにてこちら'
@@ -335,7 +336,7 @@ def test_title_in_filename(test_session, fake_now, test_directory):
assert str(archive1.screenshot_path.path.relative_to(test_directory)) == \
'archive/2000-01-01-00-00-00_ジにてこちら.png'
- def fake_request_archive(_):
+ async def fake_request_archive(_):
singlefile = '\ntest single-file\nジにてこちら\n' # no title in HTML
r = dict(
# No Title from Readability.
@@ -346,7 +347,7 @@ def test_title_in_filename(test_session, fake_now, test_directory):
return singlefile, r, s
with mock.patch('modules.archive.lib.request_archive', fake_request_archive):
- archive2 = do_archive('example.com')
+ archive2 = await do_archive('example.com')
assert str(archive2.singlefile_path.path.relative_to(test_directory)) == \
'archive/2000-01-01-00-00-00_NA.html'
@@ -359,7 +360,7 @@ def test_title_in_filename(test_session, fake_now, test_directory):
assert str(archive2.screenshot_path.path.relative_to(test_directory)) == \
'archive/2000-01-01-00-00-00_NA.png'
- def fake_request_archive(_):
+ async def fake_request_archive(_):
singlefile = '\ntest single-file\nジにてこちら\ndangerous ;\\//_title</html></html>'
r = dict(
# No Title from Readability.
@@ -370,7 +371,7 @@ def test_title_in_filename(test_session, fake_now, test_directory):
return singlefile, r, s
with mock.patch('modules.archive.lib.request_archive', fake_request_archive):
- archive3 = do_archive('example.com')
+ archive3 = await do_archive('example.com')
assert str(archive3.singlefile_path.path.relative_to(test_directory)) == \
'archive/2000-01-01-00-00-00_dangerous ;_title.html'
diff --git a/modules/videos/channel/lib.py b/modules/videos/channel/lib.py
index 98c0abe2..2dd9efbe 100644
--- a/modules/videos/channel/lib.py
+++ b/modules/videos/channel/lib.py
@@ -199,4 +199,3 @@ def download_channel(id_: int):
raise InvalidDownload(f'Channel {channel.name} does not have a download!')
download.renew(reset_attempts=True)
session.commit()
- download_manager.start_downloads()
diff --git a/modules/videos/channel/test/test_api.py b/modules/videos/channel/test/test_api.py
index 04756506..e298f9be 100644
--- a/modules/videos/channel/test/test_api.py
+++ b/modules/videos/channel/test/test_api.py
@@ -6,6 +6,7 @@ from datetime import timedelta
from http import HTTPStatus
import mock
+import pytest
from modules.videos.channel.lib import delete_channel
from modules.videos.common import get_no_channel_directory
@@ -108,7 +109,8 @@ def test_channel_no_download_frequency(test_client, test_session, test_directory
assert response.status_code == HTTPStatus.NO_CONTENT, response.json
-def test_channel_frequency_update(download_channel, test_client, test_session):
+@pytest.mark.asyncio
+def test_channel_frequency_update(download_channel, test_client, test_session, test_download_manager):
"""
A Channel's Download record is updated when the Channel's frequency is updated.
"""
@@ -336,8 +338,12 @@ def test_channel_empty_url_doesnt_conflict(test_client, test_session, test_direc
assert location != response.headers['Location']
-def test_download_channel_no_refresh(test_session, download_channel, video_download_manager):
+@pytest.mark.asyncio
+async def test_download_channel_no_refresh(test_session, download_channel, video_download_manager):
"""A Channel cannot be downloaded until it has been refreshed."""
+ d = download_channel.get_download()
+ d.next_download = video_download_manager.calculate_next_download(d)
+ test_session.commit()
def check_refreshed(expected: bool):
channel = test_session.query(Channel).one()
@@ -349,7 +355,8 @@ def test_download_channel_no_refresh(test_session, download_channel, video_downl
with mock.patch('modules.videos.downloader.YDL.extract_info') as mock_extract_info:
mock_extract_info.return_value = {'entries': [], 'url': 'foo', 'uploader': 'some uploader',
'channel_id': 'the id', 'id': 'the id'}
- video_download_manager.do_downloads_sync()
+ await video_download_manager.do_downloads()
+ await video_download_manager.wait_for_all_downloads()
check_refreshed(True)
diff --git a/modules/videos/conftest.py b/modules/videos/conftest.py
index ddc4e55f..df2c3e9c 100644
--- a/modules/videos/conftest.py
+++ b/modules/videos/conftest.py
@@ -140,9 +140,7 @@ def video_factory(test_session, test_directory):
@pytest.fixture
def video_download_manager(test_download_manager) -> DownloadManager:
- """
- Get a DownloadManager ready to download Videos and Channels.
- """
+ """Get a DownloadManager ready to download Videos and Channels."""
channel_downloader = ChannelDownloader()
video_downloader = VideoDownloader(priority=40)
diff --git a/modules/videos/downloader.py b/modules/videos/downloader.py
index cc439e55..98c95a78 100755
--- a/modules/videos/downloader.py
+++ b/modules/videos/downloader.py
@@ -81,7 +81,7 @@ class ChannelDownloader(Downloader, ABC):
# A playlist will have an id different from its channel.
return info['id'] != info['channel_id']
- def do_download(self, download: Download) -> DownloadResult:
+ async def do_download(self, download: Download) -> DownloadResult:
"""Update a Channel's catalog, then schedule downloads of every missing video."""
info = extract_info(download.url, process=False)
# Resolve the entries generator.
@@ -142,8 +142,9 @@ YT_DLP_BIN = which(
class VideoDownloader(Downloader, ABC):
- """
- Download a single video. Store the video in its channel's directory, otherwise store it in `videos/NO CHANNEL`.
+ """Downloads a single video.
+
+ Store the video in its channel's directory, otherwise store it in `videos/NO CHANNEL`.
"""
name = 'video'
pretty_name = 'Videos'
@@ -152,13 +153,9 @@ class VideoDownloader(Downloader, ABC):
return f'<VideoDownloader>'
@optional_session
- def already_downloaded(self, url: str, session: Session = None) -> Optional[Video]:
+ def already_downloaded(self, url: str, session: Session = None) -> bool:
# We only consider a video record with a video file as "downloaded".
- video = session.query(Video).filter(
- Video.url == url,
- Video.video_path != None,
- ).one_or_none()
- return video
+ return bool(session.query(Video).filter(Video.url == url, Video.video_path != None, ).count()) # noqa
@classmethod
def valid_url(cls, url) -> Tuple[bool, Optional[dict]]:
@@ -177,7 +174,7 @@ class VideoDownloader(Downloader, ABC):
logger.debug(f'{cls.__name__} not suitable for {url}')
return False, None
- def do_download(self, download: Download) -> DownloadResult:
+ async def do_download(self, download: Download) -> DownloadResult:
if download.attempts >= 10:
raise UnrecoverableDownloadError('Max download attempts reached')
@@ -226,7 +223,7 @@ class VideoDownloader(Downloader, ABC):
'--compat-options', 'no-live-chat',
url,
)
- return_code, logs = self.process_runner(url, cmd, out_dir)
+ return_code, logs = await self.process_runner(url, cmd, out_dir)
stdout = logs['stdout'].decode() if hasattr(logs['stdout'], 'decode') else logs['stdout']
stderr = logs['stderr'].decode() if hasattr(logs['stderr'], 'decode') else logs['stderr']
@@ -293,6 +290,7 @@ class VideoDownloader(Downloader, ABC):
@staticmethod
def prepare_filename(url: str, out_dir: pathlib.Path) -> Tuple[pathlib.Path, dict]:
+ """Get the full path of a video file from its URL."""
if not out_dir.is_dir():
raise ValueError(f'Output directory does not exist! {out_dir=}')
@@ -300,7 +298,6 @@ class VideoDownloader(Downloader, ABC):
options = get_downloader_config().dict()
options['outtmpl'] = f'{out_dir}/{options["file_name_format"]}'
options['merge_output_format'] = PREFERRED_VIDEO_EXTENSION
- logger.debug(f'Download video with options: {options}')
logger.debug(f'Downloading {url} to {out_dir}')
diff --git a/modules/videos/test/test_downloader.py b/modules/videos/test/test_downloader.py
index 7950f821..6ed78d2e 100644
--- a/modules/videos/test/test_downloader.py
+++ b/modules/videos/test/test_downloader.py
@@ -150,8 +150,9 @@ class TestVideosDownloaders(TestAPI):
self.assertEqual(self.mgr.get_downloader('https://www.youtube.com/watch?v=31jPEBiAC3c')[0], self.vd)
-def test_video_download_no_channel(test_session, video_download_manager, video_factory, test_directory,
- mock_video_extract_info, mock_video_process_runner):
+@pytest.mark.asyncio
+async def test_video_download_no_channel(test_session, video_download_manager, video_factory, test_directory,
+ mock_video_extract_info, mock_video_process_runner):
"""A video can be downloaded even if it does not have a Channel."""
channel_dir = test_directory / 'NO CHANNEL'
channel_dir.mkdir(parents=True)
@@ -168,7 +169,8 @@ def test_video_download_no_channel(test_session, video_download_manager, video_f
mock_video_extract_info.return_value = info_json
with mock.patch('modules.videos.downloader.VideoDownloader.prepare_filename') as mock_prepare_filename:
mock_prepare_filename.return_value = (video_path, {'id': 'foo'})
- video_download_manager.create_download(url, test_session)
+ video_download_manager.create_download(url)
+ await video_download_manager.wait_for_all_downloads()
mock_video_process_runner.assert_called_once()
@@ -176,9 +178,10 @@ def test_video_download_no_channel(test_session, video_download_manager, video_f
assert str(video.video_path.path) == f'{channel_dir}/video.mp4'
-def test_download_channel(test_session, simple_channel, video_download_manager, video_file,
- mock_video_extract_info, mock_video_prepare_filename,
- mock_video_process_runner):
+@pytest.mark.asyncio
+async def test_download_channel(test_session, simple_channel, video_download_manager, video_file,
+ mock_video_extract_info, mock_video_prepare_filename,
+ mock_video_process_runner):
"""Downloading (updating the catalog of) a Channel creates download records for all of it's missing videos."""
url = 'https://www.youtube.com/c/LearningSelfReliance/videos'
@@ -186,7 +189,8 @@ def test_download_channel(test_session, simple_channel, video_download_manager,
mock_video_extract_info.return_value = example_channel_json
with mock.patch('modules.videos.downloader.get_channel') as mock_get_channel:
mock_get_channel.return_value = simple_channel
- video_download_manager.create_download(url, test_session)
+ video_download_manager.create_download(url)
+ await video_download_manager.wait_for_all_downloads()
# Two videos are in the example channel.
downloads = video_download_manager.get_once_downloads(test_session)
@@ -238,18 +242,22 @@ def test_get_or_create_channel(test_session):
assert get_or_create_channel(source_id='quux') == channel
-def test_download_info_save(test_session, video_download_manager, mock_video_extract_info, mock_video_prepare_filename,
- mock_video_process_runner):
+@pytest.mark.asyncio
+async def test_download_info_save(test_session, video_download_manager, mock_video_extract_info,
+ mock_video_prepare_filename,
+ mock_video_process_runner):
"""
An info dict from Downloader.valid_url is saved to the Downloader and can be used by the Downloader later.
"""
mock_video_extract_info.return_value = {'some': 'info'}
mock_video_prepare_filename.return_value = 'some path'
- video_download_manager.create_download('https://www.youtube.com/watch?v=HQ_62YwcA80', test_session)
+ video_download_manager.create_download('https://www.youtube.com/watch?v=HQ_62YwcA80')
+ await video_download_manager.wait_for_all_downloads()
assert mock_video_extract_info.call_count == 2 # Only called twice. The "info" gathering call was skipped.
-def test_selected_downloader(test_session, video_download_manager, successful_download):
+@pytest.mark.asyncio
+async def test_selected_downloader(test_session, video_download_manager, successful_download):
"""
A user can specify which Downloader to use, no other should ever be used (even if the Download fails).
"""
@@ -267,7 +275,8 @@ def test_selected_downloader(test_session, video_download_manager, successful_do
# DownloadManager will attempt to choose the Downloader.
with mock.patch('modules.videos.downloader.YDL.extract_info') as mock_extract_info:
mock_extract_info.side_effect = UnsupportedError('not this one')
- video_download_manager.create_download('https://example.com', test_session)
+ video_download_manager.create_download('https://example.com')
+ await video_download_manager.wait_for_all_downloads()
mock_http_valid_url.assert_called_once()
mock_extract_info.assert_called_once()
mock_do_download.assert_called_once()
@@ -279,7 +288,8 @@ def test_selected_downloader(test_session, video_download_manager, successful_do
# DownloadManager will only use the HTTPDownloader.
with mock.patch('modules.videos.downloader.YDL.extract_info') as mock_extract_info:
- video_download_manager.create_download('https://example.com', test_session, downloader='http')
+ video_download_manager.create_download('https://example.com', downloader='http')
+ await video_download_manager.wait_for_all_downloads()
# DownloadManager does not try to find the Download, it trusts the "downloader" param.
mock_extract_info.assert_not_called()
mock_http_valid_url.assert_not_called()
@@ -294,7 +304,8 @@ def test_selected_downloader(test_session, video_download_manager, successful_do
# DownloadManager will only use the HTTPDownloader, even if the download fails.
with mock.patch('modules.videos.downloader.YDL.extract_info') as mock_extract_info:
mock_do_download.side_effect = Exception('oh no!')
- video_download_manager.create_download('https://example.com', test_session, downloader='http')
+ video_download_manager.create_download('https://example.com', downloader='http')
+ await video_download_manager.wait_for_all_downloads()
# DownloadManager does not try to find the Download, it trusts the "downloader" param.
mock_extract_info.assert_not_called()
mock_http_valid_url.assert_not_called()
@@ -312,12 +323,13 @@ def test_channel_downloader_hidden(video_download_manager):
assert [i.__json__() for i in downloaders] == [(dict(name='video', pretty_name='Videos')), ]
+@pytest.mark.asyncio
def test_bad_downloader(test_session, video_download_manager):
"""
Attempting to use an unknown downloader should raise an error.
"""
with pytest.raises(InvalidDownload):
- video_download_manager.create_download('https://example.com', test_session, 'bad downloader')
+ video_download_manager.create_download('https://example.com', downloader='bad downloader')
def test_channel_download_no_download(test_session, video_download_manager, simple_channel):
@@ -332,22 +344,27 @@ def test_channel_download_no_download(test_session, video_download_manager, simp
download_channel(simple_channel.id)
-def test_invalid_download_url(test_session, test_download_manager, mock_video_extract_info):
+@pytest.mark.asyncio
+async def test_invalid_download_url(test_session, test_download_manager, mock_video_extract_info):
"""An invalid url should not be attempted again."""
_, session = get_db_context()
with mock.patch('modules.videos.downloader.VideoDownloader.valid_url') as mock_valid_url, \
mock.patch('wrolpi.downloader.DownloadManager.get_downloader') as mock_get_downloader:
- mock_get_downloader.return_value = (VideoDownloader(), None)
+ video_downloader = VideoDownloader()
+ test_download_manager.register_downloader(video_downloader)
+ mock_get_downloader.return_value = (video_downloader, None)
mock_valid_url.return_value = (True, {}) # url is valid, but is not a video URL.
mock_video_extract_info.side_effect = UnsupportedError('oops')
- test_download_manager.create_download('invalid url', session)
+ test_download_manager.create_download('invalid url', downloader=VideoDownloader.name)
+ await test_download_manager.wait_for_all_downloads()
download = session.query(Download).one()
assert download.status == 'failed'
-def test_video_download_1(test_session, test_directory, simple_channel, video_download_manager,
- mock_video_process_runner):
+@pytest.mark.asyncio
+async def test_video_download_1(test_session, test_directory, simple_channel, video_download_manager,
+ mock_video_process_runner):
simple_channel.source_id = example_video_json['channel_id']
simple_channel.directory = test_directory / 'videos/channel name'
simple_channel.directory.mkdir(parents=True)
@@ -361,7 +378,8 @@ def test_video_download_1(test_session, test_directory, simple_channel, video_do
mock_extract_info.return_value = example_video_json
mock_prepare_filename.return_value = (video_file, {'id': 'foo'})
- video_download_manager.create_download(url, test_session)
+ video_download_manager.create_download(url)
+ await video_download_manager.wait_for_all_downloads()
mock_video_process_runner.assert_called_once()
video_url, _, out_dir = mock_video_process_runner.call_args[0]
@@ -374,8 +392,9 @@ def test_video_download_1(test_session, test_directory, simple_channel, video_do
assert video.video_path.path.is_absolute()
-def test_download_result(test_session, test_directory, video_download_manager, mock_video_process_runner,
- mock_video_extract_info):
+@pytest.mark.asyncio
+async def test_download_result(test_session, test_directory, video_download_manager, mock_video_process_runner,
+ mock_video_extract_info):
"""VideoDownloader returns a DownloadResult when complete."""
channel_directory = test_directory / 'videos/channel name'
channel_directory.mkdir(parents=True)
@@ -385,7 +404,8 @@ def test_download_result(test_session, test_directory, video_download_manager, m
mock_video_extract_info.return_value = example_video_json
with mock.patch('modules.videos.downloader.VideoDownloader.prepare_filename') as mock_prepare_filename:
mock_prepare_filename.return_value = [video_file, {'id': 'foo'}]
- video_download_manager.create_download('https://example.com', test_session)
+ video_download_manager.create_download('https://example.com')
+ await video_download_manager.wait_for_all_downloads()
download: Download = test_session.query(Download).one()
assert download.url == 'https://example.com'
@@ -451,7 +471,8 @@ example_playlist_json = {
}
-def test_download_playlist(test_session, test_directory, mock_video_extract_info, video_download_manager):
+@pytest.mark.asyncio
+async def test_download_playlist(test_session, test_directory, mock_video_extract_info, video_download_manager):
"""
All videos in a playlist can be downloaded for it's Channel.
"""
@@ -467,7 +488,7 @@ def test_download_playlist(test_session, test_directory, mock_video_extract_info
with mock.patch('modules.videos.downloader.VideoDownloader.do_download') as mock_video_do_download:
mock_video_do_download.return_value = DownloadResult(success=True) # Don't download the videos.
- result = channel_downloader.do_download(download)
+ result = await channel_downloader.do_download(download)
assert result.success is True, 'Download was not successful'
assert set(result.downloads) == {
'https://youtube.com/watch?v=video_2_url',
diff --git a/requirements.txt b/requirements.txt
index 109d6658..06ab2f28 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,23 +1,23 @@
+Pillow==9.0.1
+Pint==0.18
+PyYAML==5.3.1 # this is the only version available on RPi
+SQLAlchemy==1.3.22
+aiohttp==3.8.1
alembic==1.7.5
beautifulsoup4==4.10.0
cachetools==5.0.0
feedparser==6.0.9
mock==4.0.3
-Pillow==9.0.1
-Pint==0.18
psycopg2-binary==2.9.3
pytest-asyncio==0.17.2
pytest==6.2.5
python-dotenv==0.19.2
pytz==2021.3
-PyYAML==5.3.1 # this is the only version available on RPi
-requests==2.27.1
sanic-ext==22.1.2
sanic-testing==0.8.2
sanic==21.12.1
selenium==4.1.0
-SQLAlchemy==1.3.22
srt==3.5.1
websockets==10.1
webvtt-py==0.4.6
-yt-dlp==2022.6.29
+yt-dlp==2022.7.18
diff --git a/wrolpi/common.py b/wrolpi/common.py
index b0e40cd0..317bd3f8 100644
--- a/wrolpi/common.py
+++ b/wrolpi/common.py
@@ -20,6 +20,7 @@ from pathlib import Path
from typing import Union, Callable, Tuple, Dict, List, Iterable, Optional, Generator, Any
from urllib.parse import urlunsplit, urlparse
+import aiohttp
import yaml
from sanic import Blueprint
from sanic.request import Request
@@ -370,9 +371,7 @@ def set_wrol_mode(enable: bool):
def wrol_mode_check(func):
- """
- Wraps a function so that it cannot be called when WROL Mode is enabled.
- """
+ """Wraps a function so that it cannot be called when WROL Mode is enabled."""
@wraps(func)
def check(*a, **kw):
@@ -786,3 +785,11 @@ def any_extensions(filename: str, extensions: Iterable = ()):
Matches lower or upper case of the extension.
"""
return any(filename.lower().endswith(ext) for ext in extensions)
+
+
+async def aiohttp_post(url: str, json_, timeout: int = None) -> Tuple[Dict, int]:
+ """Perform an async aiohttp post request. Return the json contents."""
+ timeout = aiohttp.ClientTimeout(total=timeout) if timeout else None
+ async with aiohttp.ClientSession(timeout=timeout) as session:
+ async with session.post(url, json=json_) as response:
+ return await response.json(), response.status
diff --git a/wrolpi/conftest.py b/wrolpi/conftest.py
index 47cea0e1..70b5ecdc 100644
--- a/wrolpi/conftest.py
+++ b/wrolpi/conftest.py
@@ -1,6 +1,7 @@
"""
Fixtures for Pytest tests.
"""
+import asyncio
import pathlib
import tempfile
from typing import Tuple, Set
@@ -121,9 +122,20 @@ def test_download_manager_config(test_directory):
@pytest.fixture
-def test_download_manager(test_download_manager_config):
+async def test_download_manager(
+ test_session, # session is required because downloads can start without the test DB in place.
+ test_download_manager_config,
+):
manager = DownloadManager()
- return manager
+ try:
+ loop = asyncio.get_running_loop()
+ except RuntimeError:
+ loop = asyncio.get_event_loop()
+ manager.start_workers(loop=loop)
+
+ yield manager
+
+ manager.cancel_workers()
@pytest.fixture
diff --git a/wrolpi/downloader.py b/wrolpi/downloader.py
index 61950a76..4cd1581b 100644
--- a/wrolpi/downloader.py
+++ b/wrolpi/downloader.py
@@ -1,16 +1,17 @@
import asyncio
import multiprocessing
import pathlib
-import subprocess
import traceback
from abc import ABC
+from asyncio import Queue, Task, QueueEmpty
from dataclasses import dataclass, field
from datetime import timedelta, datetime
from enum import Enum
from functools import partial
from operator import attrgetter
-from typing import List, Dict
+from typing import List, Dict, Generator
from typing import Tuple, Optional
+from urllib.parse import urlparse
import feedparser
from feedparser import FeedParserDict
@@ -20,15 +21,12 @@ from sqlalchemy.orm import Session
from wrolpi.common import Base, ModelHelper, logger, wrol_mode_check, zig_zag, ConfigFile, WROLPI_CONFIG
from wrolpi.dates import TZDateTime, now, Seconds, local_timezone, recursive_replace_tz
-from wrolpi.db import get_db_session, get_db_curs, get_db_context, optional_session
+from wrolpi.db import get_db_session, get_db_curs, optional_session
from wrolpi.errors import InvalidDownload, UnrecoverableDownloadError
from wrolpi.vars import PYTEST
logger = logger.getChild(__name__)
-DOWNLOAD_IN_PROGRESS = multiprocessing.Semaphore()
-DEFAULT_RETRY_FREQUENCY = timedelta(hours=1)
-
class DownloadFrequency(int, Enum):
hourly = 3600
@@ -52,7 +50,7 @@ class DownloadResult:
class Download(ModelHelper, Base):
"""Model that is used to schedule downloads."""
- __tablename__ = 'download'
+ __tablename__ = 'download' # noqa
id = Column(Integer, primary_key=True)
url = Column(String, nullable=False)
@@ -66,19 +64,18 @@ class Download(ModelHelper, Base):
next_download = Column(TZDateTime)
status = Column(String, default='new')
sub_downloader = Column(Text)
+ _manager = None
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
- self.manager: DownloadManager = None
-
def __repr__(self):
if self.next_download or self.frequency:
return f'<Download id={self.id} status={self.status} url={repr(self.url)} ' \
- f'next_download={self.next_download} frequency={self.frequency} attempts={self.attempts} ' \
- f'error={self.error}>'
+ f'next_download={repr(self.next_download)} frequency={self.frequency} attempts={self.attempts} ' \
+ f'error={bool(self.error)}>'
return f'<Download id={self.id} status={self.status} url={repr(self.url)} attempts={self.attempts} ' \
- f'error={self.error}>'
+ f'error={bool(self.error)}>'
def __json__(self):
d = dict(
@@ -126,6 +123,20 @@ class Download(ModelHelper, Base):
return self.manager.get_downloader(self.url)
+ @property
+ def domain(self):
+ return urlparse(self.url).netloc
+
+ @property
+ def manager(self):
+ if self._manager:
+ return self._manager
+ raise ValueError('No manager has been set!')
+
+ @manager.setter
+ def manager(self, value):
+ self._manager = value
+
class Downloader:
name: str = None
@@ -156,11 +167,14 @@ class Downloader:
def __json__(self):
return dict(name=self.name, pretty_name=self.pretty_name)
+ def __repr__(self):
+ return f'<Downloader name={self.name}>'
+
@classmethod
def valid_url(cls, url: str) -> Tuple[bool, Optional[dict]]:
raise NotImplementedError()
- def do_download(self, download: Download) -> DownloadResult:
+ async def do_download(self, download: Download) -> DownloadResult:
raise NotImplementedError()
@optional_session
@@ -187,8 +201,8 @@ class Downloader:
if self._kill.is_set():
self._kill.clear()
- def process_runner(self, url: str, cmd: Tuple[str, ...], cwd: pathlib.Path, timeout: int = None,
- **kwargs) -> Tuple[int, dict]:
+ async def process_runner(self, url: str, cmd: Tuple[str, ...], cwd: pathlib.Path, timeout: int = None,
+ **kwargs) -> Tuple[int, dict]:
"""
Run a subprocess using the provided arguments. This process can be killed by the Download Manager.
@@ -196,7 +210,11 @@ class Downloader:
"""
logger.debug(f'{self} launching download process with args: {" ".join(cmd)}')
start = now()
- proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, cwd=cwd, **kwargs)
+ proc = await asyncio.create_subprocess_exec(*cmd,
+ stdout=asyncio.subprocess.PIPE,
+ stderr=asyncio.subprocess.PIPE,
+ cwd=cwd,
+ **kwargs)
pid = proc.pid
# Timeout can be any positive integer. Global download timeout takes precedence, unless it is 0.
@@ -204,37 +222,53 @@ class Downloader:
timeout = WROLPI_CONFIG.download_timeout or timeout or self.timeout
logger.debug(f'{self} launched download process {pid=} {timeout=} for {url}')
+ stdout, stderr = None, None
try:
while True:
- try:
- proc.wait(timeout=1.0)
- # Download finished.
+ task = asyncio.Task(proc.communicate())
+ done, pending = await asyncio.wait([task, ], timeout=1)
+ # Cancel communicate.
+ if pending:
+ pending.pop().cancel()
+ if done:
+ # Process finished, get the result of proc.communicate().
+ stdout, stderr = done.pop().result()
+ break
+ if proc.returncode:
+ # Process died.
+ stdout, stderr = await proc.communicate()
break
- except subprocess.TimeoutExpired:
- pass
- if timeout and (elapsed := now() - start).total_seconds() > timeout:
+ elapsed = (now() - start).total_seconds()
+ if timeout and elapsed > timeout:
logger.warning(f'Download has exceeded its timeout {elapsed=}')
self.kill()
if self._kill.is_set():
- logger.warning(f'Killing download {pid=}')
+ logger.warning(f'Killing download {pid=}, {elapsed} seconds elapsed (timeout was not exceeded).')
proc.kill()
- proc.poll()
break
+ except Exception as e:
+ logger.error(f'{self}.process_runner had a download error', exc_info=e)
+ raise
finally:
self.clear()
# Output all logs from the process.
# TODO is there a way to stream this output while the process is running?
logger.debug(f'Download exited with {proc.returncode}')
- stdout, stderr = proc.communicate()
logs = {'stdout': stdout, 'stderr': stderr}
return proc.returncode, logs
class DownloadManager:
+ """
+ Runs a collection of workers which will download any URLs in the `download` table.
+
+ Only one download from each domain will be downloaded at a time. If we only have one domain (example.com) in the
+ URLs to be downloaded, then only one worker will be busy.
+ """
priority_sorter = partial(sorted, key=attrgetter('priority'))
def __init__(self):
@@ -242,6 +276,12 @@ class DownloadManager:
self._instances = dict()
self.disabled = multiprocessing.Event()
self.stopped: bool = False
+ self.download_queue: Queue = Queue()
+ self.workers: List[Task] = []
+ self.worker_count: int = 4
+ self.data = multiprocessing.Manager().dict()
+ # We haven't started downloads yet, so no domains are downloading.
+ self.data['processing_domains'] = []
def register_downloader(self, instance: Downloader):
if not isinstance(instance, Downloader):
@@ -255,6 +295,132 @@ class DownloadManager:
self.instances = tuple(self.priority_sorter(i))
self._instances[instance.name] = instance
+ async def download_worker(self, num: int, queue: Queue):
+ """Fetch a download from the queue, perform the download then store the results.
+
+ Calls DownloadManger.start_downloads() after a download completes.
+ """
+ logger.debug(f'DownloadManager.download_worker {num} starting up')
+ from wrolpi.db import get_db_session
+
+ while True:
+ if self.stopped:
+ return
+
+ try:
+ download: Download = queue.get_nowait()
+ download_id, url = download.id, download.url
+ logger.debug(f'Download worker {num} got download {download}')
+
+ downloader: Downloader = download.get_downloader()
+ if not downloader:
+ logger.warning(f'Could not find downloader for {download.downloader=}')
+
+ with get_db_session(commit=True) as session:
+ # Mark the download as started in new session so the change is committed.
+ download = session.query(Download).filter_by(id=download_id).one()
+ download.started()
+
+ self.data['processing_domains'].append(download.domain)
+
+ try_again = True
+ try:
+ if asyncio.iscoroutinefunction(downloader.do_download):
+ result = await downloader.do_download(download)
+ else:
+ result = downloader.do_download(download)
+ except UnrecoverableDownloadError as e:
+ # Download failed and should not be retried.
+ logger.warning(f'UnrecoverableDownloadError for {url}', exc_info=e)
+ result = DownloadResult(success=False, error=str(traceback.format_exc()))
+ try_again = False
+ except Exception as e:
+ logger.warning(f'Failed to download {url}. Will be tried again later.', exc_info=e)
+ result = DownloadResult(success=False, error=str(traceback.format_exc()))
+
+ error_len = len(result.error) if result.error else 0
+ logger.debug(f'Got success={result.success} from {downloader} with {error_len=}')
+
+ with get_db_session(commit=True) as session:
+ # Modify the download in a new session because downloads may take a long time.
+ download = session.query(Download).filter_by(id=download_id).one()
+ # Use a new location if provided, keep the old location if no new location is provided, otherwise
+ # clear out an outdated location.
+ download.location = result.location or download.location or None
+ # Clear any old errors if the download succeeded.
+ download.error = result.error if result.error else None
+ download.next_download = self.calculate_next_download(download, session)
+
+ if result.downloads:
+ logger.info(f'Adding {len(result.downloads)} downloads from result of {download.url}')
+ self.create_downloads(result.downloads, session,
+ downloader=download.sub_downloader)
+
+ if try_again is False and not download.frequency:
+ # Only once-downloads can fail.
+ download.fail()
+ elif result.success:
+ download.complete()
+ else:
+ download.defer()
+
+ queue.task_done()
+ # Remove this domain from the running list.
+ self._remove_domain(download.domain)
+ # Request any new downloads be added to the queue.
+ asyncio.create_task(self.queue_downloads())
+ except asyncio.CancelledError:
+ logger.warning(f'download_worker canceled!')
+ queue.task_done()
+ return
+ except QueueEmpty:
+ # No work yet.
+ await asyncio.sleep(0.1)
+ continue
+ except Exception as e:
+ logger.warning(f'Download worker had unexpected error', exc_info=e)
+
+ def _add_domain(self, domain: str):
+ """Add a domain to the processing list.
+
+ Raises ValueError if the domain is already processing."""
+ processing_domains = self.data['processing_domains']
+ if domain in processing_domains:
+ raise ValueError(f'Domain already being downloaded! {domain}')
+ self.data['processing_domains'] = [*processing_domains, domain]
+
+ def _remove_domain(self, domain: str):
+ """Remove a domain from the processing list."""
+ self.data['processing_domains'] = [i for i in self.data['processing_domains'] if i != domain]
+
+ @wrol_mode_check
+ def start_workers(self, loop=None):
+ """Start all download worker tasks. Does nothing if they are already running."""
+ if not self.workers:
+ if not loop:
+ try:
+ loop = asyncio.get_running_loop()
+ except RuntimeError:
+ # Loop isn't running, start one. Probably testing.
+ loop = asyncio.get_event_loop()
+
+ logger.debug(f'Starting {self.worker_count} download workers')
+ for i in range(self.worker_count):
+ coro = self.download_worker(i, self.download_queue)
+ task = loop.create_task(coro)
+ self.workers.append(task)
+
+ def workers_running(self):
+ for task in self.workers:
+ if not task.done():
+ return True
+ return False
+
+ def cancel_workers(self):
+ if self.workers_running():
+ for task in self.workers:
+ task.cancel()
+
def get_downloader(self, url: str) -> Tuple[Downloader, Dict]:
for i in self.instances:
valid, info = i.valid_url(url)
@@ -262,29 +428,10 @@ class DownloadManager:
return i, info
raise InvalidDownload(f'Invalid URL {url=}')
- def get_downloader_by_name(self, name: str) -> Downloader:
- """
- Tries to find a registered Downloader by its name. Returns None if it cannot be found.
- """
+ def get_downloader_by_name(self, name: str) -> Optional[Downloader]:
+ """Attempt to find a registered Downloader by its name. Returns None if it cannot be found."""
return self._instances.get(name)
- def get_new_downloads(self, session: Session) -> List[Download]:
- """
- Get all "new" downloads. This method fetches the first download each iteration, so it will fetch downloads
- that were created after calling it.
- """
- last = None
- while True:
- download = session.query(Download).filter_by(status='new').order_by(Download.id).first()
- if not download:
- return
- if download == last:
- # Got the last download again. Is something wrong?
- return
- last = download
- download.manager = self
- yield download
-
def get_or_create_download(self, url: str, session: Session) -> Download:
"""Get a Download by its URL, if it cannot be found create one."""
if not url:
@@ -301,159 +448,136 @@ class DownloadManager:
download.manager = self
return download
- def create_download(self, url: str, session, downloader: str = None, skip_download: bool = False,
- reset_attempts: bool = False, sub_downloader: str = None) -> Download:
+ @optional_session
+ def create_download(self, url: str, session: Session = None, downloader: str = None, reset_attempts: bool = False,
+ sub_downloader: str = None) -> Download:
"""Schedule a URL for download. If the URL failed previously, it may be retried."""
- downloads = self.create_downloads([url], session, downloader, skip_download=skip_download,
+ downloads = self.create_downloads([url], session=session, downloader=downloader,
reset_attempts=reset_attempts, sub_downloader=sub_downloader)
return downloads[0]
- def create_downloads(self, urls: List[str], session: Session = None, downloader_name: str = None,
- skip_download: bool = False, reset_attempts: bool = False, sub_downloader: str = None) \
+ @optional_session
+ def create_downloads(self, urls: List[str], session: Session = None, downloader: str = None,
+ reset_attempts: bool = False, sub_downloader: str = None) \
-> List[Download]:
"""Schedule all URLs for download. If one cannot be downloaded, none will be added."""
- if not session:
- _, session = get_db_context()
+ if not all(urls):
+ raise ValueError('Download must have a URL')
downloads = []
- downloader = self.get_downloader_by_name(downloader_name) if downloader_name else None
- if downloader_name and not downloader:
- # Could not find the downloader
- raise InvalidDownload(f'Unknown downloader {downloader_name}')
+ shared_downloader = self.get_downloader_by_name(downloader) if downloader else None
+ if downloader and not shared_downloader:
+ # Could not look up the downloader.
+ raise InvalidDownload(f'Unknown downloader {downloader}')
with session.transaction:
for url in urls:
- if not url:
- raise ValueError('Download must have a URL')
if url in DOWNLOAD_MANAGER_CONFIG.skip_urls and reset_attempts:
# User manually entered this download, remove it from the skip list.
- DOWNLOAD_MANAGER_CONFIG.skip_urls = [i for i in DOWNLOAD_MANAGER_CONFIG.skip_urls if i != url]
- DOWNLOAD_MANAGER_CONFIG.save()
+ self.remove_from_skip_list(url)
elif url in DOWNLOAD_MANAGER_CONFIG.skip_urls:
logger.warning(f'Skipping {url} because it is in the download_manager.yaml skip list.')
continue
info_json = None
- if not downloader:
+ if not shared_downloader:
# User has requested automatic downloader selection, try and find it.
- downloader_, info_json = self.get_downloader(url)
+ local_downloader, info_json = self.get_downloader(url)
else:
- downloader_ = downloader
+ # One Downloader provided for all URLs, use it.
+ local_downloader = shared_downloader
download = self.get_or_create_download(url, session)
# Download may have failed, try again.
download.renew(reset_attempts=reset_attempts)
- download.downloader = downloader_.name
+ download.downloader = local_downloader.name
download.info_json = info_json or download.info_json
download.sub_downloader = sub_downloader
downloads.append(download)
- if skip_download is True:
- return downloads
-
- if PYTEST:
- # Download now (for testing).
- self._do_downloads(session)
- else:
- # Start downloading ASAP.
- self.start_downloads()
+ # Start downloading ASAP.
+ try:
+ asyncio.create_task(self.queue_downloads())
+ except RuntimeError:
+ # Event loop isn't running. Probably testing?
+ if not PYTEST:
+ logger.info(f'Unable to queue downloads after creating {len(downloads)} download(s).')
return downloads
+ def recurring_download(self, url: str, frequency: int, downloader: str = None,
+ sub_downloader: str = None, reset_attempts: bool = False) -> Download:
+ """Schedule a recurring download."""
+ if not frequency:
+ raise ValueError('Recurring download must have a frequency!')
+
+ with get_db_session(commit=True) as session:
+ download, = self.create_downloads([url, ], session=session, downloader=downloader,
+ sub_downloader=sub_downloader, reset_attempts=reset_attempts)
+ download.frequency = frequency
+
+ return download
+
@wrol_mode_check
@optional_session
- def _do_downloads(self, session: Session = None):
- """This method calls the Downloader's do_download method."""
+ async def queue_downloads(self, session: Session = None):
+ """Put all downloads in queue. Will only queue downloads if there are workers to take them. Each worker
+ only receives one domain, this is to prevent downloading from one domain many times at once."""
if self.disabled.is_set():
raise InvalidDownload('DownloadManager is disabled')
- try:
- downloads = self.get_new_downloads(session)
-
- download_count = 0
- for download in downloads:
- if self.disabled.is_set():
- raise InvalidDownload('DownloadManager is disabled')
- if self.stopped:
- return
-
- download_count += 1
- download_id = download.id
- url = download.url
-
- if url in DOWNLOAD_MANAGER_CONFIG.skip_urls:
- logger.info(f'Skipping download in skip_urls list: {url}')
- session.query(Download).filter_by(id=download.id).delete()
- session.commit()
- continue
-
- downloader = self.get_downloader_by_name(download.downloader)
- if not downloader:
- downloader, info_json = self.get_downloader(download.url)
- download.downloader = downloader.name
- download.info_json = info_json
- downloader.clear()
-
- download.started()
- session.commit()
-
- try_again = True
- try:
- result = downloader.do_download(download)
- except UnrecoverableDownloadError as e:
- # Download failed and should not be retried.
- logger.warning(f'UnrecoverableDownloadError for {url}', exc_info=e)
- result = DownloadResult(success=False, error=str(traceback.format_exc()))
- try_again = False
- except Exception as e:
- logger.warning(f'Failed to download {url}. Will be tried again later.', exc_info=e)
- result = DownloadResult(success=False, error=str(traceback.format_exc()))
-
- error_len = len(result.error) if result.error else 0
- logger.debug(f'Got success={result.success} from {downloader} with {error_len=}')
-
- # Long-running, get the download again.
- with get_db_session(commit=True) as session_:
- download = self.get_download(session_, id_=download_id)
- # Use a new location if provided, keep the old location if no new location is provided, otherwise
- # clear out an outdated location.
- download.location = result.location or download.location or None
- # Clear any old errors if the download succeeded.
- download.error = result.error if result.error else None
- download.next_download = self.calculate_next_download(download, session)
-
- if result.downloads:
- logger.info(f'Adding {len(result.downloads)} downloads from result of {download.url}')
- self.create_downloads(result.downloads, session_, downloader_name=download.sub_downloader)
-
- if try_again is False:
- download.fail()
- elif result.success:
- download.complete()
- else:
- download.defer()
-
- if download_count:
- logger.info(f'Done doing {download_count} downloads.')
- finally:
- session.close()
-
- def do_downloads_sync(self):
- acquired = DOWNLOAD_IN_PROGRESS.acquire(block=False)
- if not acquired:
+ if len(self.data['processing_domains']) >= len(self.workers):
return
- try:
- self._do_downloads()
- finally:
- DOWNLOAD_IN_PROGRESS.release()
+ # Find download whose domain isn't already being downloaded.
+ new_downloads = session.query(Download).filter(
+ Download.status == 'new',
+ Download.domain not in self.data['processing_domains'],
+ ).order_by(Download.id) # noqa
+ count = 0
+ for download in new_downloads:
+ download.manager = self # Assign this Download to this manager.
+ domain = download.domain
+ if domain not in self.data['processing_domains']:
+ self._add_domain(domain)
+ await self.download_queue.put(download)
+ count += 1
+ if count:
+ logger.debug(f'Added {count} downloads to queue.')
async def do_downloads(self):
- return self.do_downloads_sync()
+ """Schedule any downloads that are new.
- def start_downloads(self):
- """Start an async task to do downloads. This does nothing when testing!"""
- if not PYTEST:
- asyncio.create_task(self.do_downloads())
+ Warning: Downloads will still be running even after this returns! See `wait_for_all_downloads`.
+ """
+ if self.disabled.is_set():
+ return
+ self.renew_recurring_downloads()
+ self.delete_old_once_downloads()
+ await self.queue_downloads()
+
+ async def wait_for_all_downloads(self):
+ """Wait for all Downloads in queue AND any new Downloads to complete.
+
+ THIS METHOD IS FOR TESTING.
+ """
+ while True:
+ await asyncio.sleep(0.1)
+
+ if not self.workers_running():
+ raise ValueError('No workers are running!')
+
+ await self.queue_downloads()
+
+ try:
+ next(self.get_new_downloads())
+ continue
+ except StopIteration:
+ pass
+
+ if self.download_queue.empty():
+ # Queue is empty.
+ break
+ await self.download_queue.join()
@staticmethod
def reset_downloads():
@@ -461,23 +585,26 @@ class DownloadManager:
with get_db_curs(commit=True) as curs:
curs.execute("UPDATE download SET status='new' WHERE status='pending' OR status='deferred'")
- def recurring_download(self, url: str, frequency: int, skip_download: bool = False, downloader: str = None,
- sub_downloader: str = None, reset_attempts: bool = False) -> Download:
- """Schedule a recurring download."""
- with get_db_session(commit=True) as session:
- download = self.create_download(url, session=session, skip_download=True, downloader=downloader,
- sub_downloader=sub_downloader, reset_attempts=reset_attempts)
- download.frequency = frequency
-
- if skip_download is False and PYTEST:
- self.do_downloads_sync()
- elif skip_download is False:
- self.start_downloads()
-
- return download
-
DOWNLOAD_SORT = ('pending', 'failed', 'new', 'deferred', 'complete')
+ @optional_session
+ def get_new_downloads(self, session: Session) -> Generator[Download, None, None]:
+ """
+ Get all "new" downloads. This method fetches the first download each iteration, so it will fetch downloads
+ that were created after calling it.
+ """
+ last = None
+ while True:
+ download = session.query(Download).filter_by(status='new').order_by(Download.id).first()
+ if not download:
+ return
+ if download == last:
+ # Got the last download again. Is something wrong?
+ return
+ last = download
+ download.manager = self
+ yield download
+
@optional_session
def get_recurring_downloads(self, session: Session = None, limit: int = None):
"""Get all Downloads that will be downloaded in the future."""
@@ -519,11 +646,7 @@ class DownloadManager:
renewed = True
if renewed:
- logger.debug('Renewed recurring downloads')
session.commit()
- self.start_downloads()
- else:
- logger.debug('No recurring downloads to renew')
def get_downloads(self, session: Session) -> List[Download]:
downloads = list(session.query(Download).all())
@@ -566,11 +689,13 @@ class DownloadManager:
def stop(self):
"""Stop all Downloaders, defer all pending downloads."""
+ logger.warning('Stopping all downloads')
self.stopped = True
for downloader in self.instances:
downloader.kill()
for download in self.get_pending_downloads():
download.defer()
+ self.cancel_workers()
def kill(self):
"""Kill all downloads. Do not start new downloads."""
@@ -582,15 +707,19 @@ class DownloadManager:
def enable(self):
"""Enable downloading. Start downloading."""
+ logger.info('Enabling downloading')
+ for downloader in self.instances:
+ downloader.clear()
self.disabled.clear()
- self.start_downloads()
+ self.start_workers()
+ asyncio.create_task(self.do_downloads())
FINISHED_STATUSES = ('complete', 'failed')
def delete_old_once_downloads(self):
- """
- Delete all once-downloads that have expired. Do not delete downloads that are new, or should be tried again.
- """
+ """Delete all once-downloads that have expired.
+
+ Do not delete downloads that are new, or should be tried again."""
with get_db_session(commit=True) as session:
downloads = self.get_once_downloads(session)
one_month = now() - timedelta(days=30)
@@ -674,8 +803,9 @@ class DownloadManager:
download.manager = self
return downloads
+ @staticmethod
@optional_session
- def calculate_next_download(self, download: Download, session: Session) -> Optional[datetime]:
+ def calculate_next_download(download: Download, session: Session) -> Optional[datetime]:
"""
If the download is "deferred", download soon. But, slowly increase the time between attempts.
@@ -738,13 +868,22 @@ class DownloadManager:
# Add all downloads to permanent skip list.
ids = [i.id for i in failed_downloads]
- urls = list(set(DOWNLOAD_MANAGER_CONFIG.skip_urls) | {i.url for i in failed_downloads})
- DOWNLOAD_MANAGER_CONFIG.skip_urls = urls
+ self.add_to_skip_list(*(i.url for i in failed_downloads))
# Delete all failed once-downloads.
session.execute('DELETE FROM download WHERE id = ANY(:ids)', {'ids': ids})
session.commit()
+ @staticmethod
+ def add_to_skip_list(*urls: str):
+ DOWNLOAD_MANAGER_CONFIG.skip_urls = list(set(DOWNLOAD_MANAGER_CONFIG.skip_urls) | set(urls))
+ DOWNLOAD_MANAGER_CONFIG.save()
+
+ @staticmethod
+ def remove_from_skip_list(url: str):
+ DOWNLOAD_MANAGER_CONFIG.skip_urls = [i for i in DOWNLOAD_MANAGER_CONFIG.skip_urls if i != url]
+ DOWNLOAD_MANAGER_CONFIG.save()
+
# The global DownloadManager. This should be used everywhere!
download_manager = DownloadManager()
@@ -824,14 +963,12 @@ class RSSDownloader(Downloader, ABC):
logger.warning(f'RSS entry {idx} did not have a link!')
session = Session.object_session(download)
- self.manager.create_downloads(urls, session=session, downloader_name=download.sub_downloader)
+ self.manager.create_downloads(urls, session=session, downloader=download.sub_downloader)
return DownloadResult(success=True)
@staticmethod
def acceptable_bozo_errors(feed):
- """
- Feedparser can report some errors, some we can ignore.
- """
+ """Feedparser can report some errors, some we can ignore."""
if 'document declared as' in str(feed['bozo_exception']):
# Feed's encoding does not match what is declared, this is fine.
return True
diff --git a/wrolpi/root_api.py b/wrolpi/root_api.py
index d481a1fc..f672e0a3 100644
--- a/wrolpi/root_api.py
+++ b/wrolpi/root_api.py
@@ -13,7 +13,6 @@ from sanic import Sanic, response, Blueprint, __version__ as sanic_version
from sanic.blueprint_group import BlueprintGroup
from sanic.request import Request
from sanic.response import HTTPResponse
-from sanic.signals import Event
from sanic_ext import validate
from sanic_ext.extensions.openapi import openapi
@@ -26,7 +25,7 @@ from wrolpi.downloader import download_manager
from wrolpi.errors import WROLModeEnabled, InvalidTimezone, API_ERRORS, APIError, ValidationError, HotspotError
from wrolpi.media_path import MediaPath
from wrolpi.schema import RegexRequest, RegexResponse, SettingsRequest, SettingsResponse, DownloadRequest, EchoResponse
-from wrolpi.vars import DOCKERIZED, PYTEST
+from wrolpi.vars import DOCKERIZED
from wrolpi.version import __version__
logger = logger.getChild(__name__)
@@ -210,7 +209,7 @@ async def post_download(_: Request, body: DownloadRequest):
download_manager.recurring_download(urls[0], body.frequency, downloader=downloader,
sub_downloader=body.sub_downloader, reset_attempts=True)
else:
- download_manager.create_downloads(urls, downloader_name=downloader, sub_downloader=body.sub_downloader,
+ download_manager.create_downloads(urls, downloader=downloader, sub_downloader=body.sub_downloader,
reset_attempts=True)
return response.empty()
@@ -376,10 +375,3 @@ def json_error_handler(request, exception: Exception):
api_app.error_handler.add(APIError, json_error_handler)
-
-
-@api_app.signal(Event.SERVER_SHUTDOWN_BEFORE)
-def handle_server_shutdown(*args, **kwargs):
- """Stop downloads when server is shutting down."""
- if not PYTEST:
- download_manager.stop()
diff --git a/wrolpi/test/common.py b/wrolpi/test/common.py
index 344ee401..6fa03feb 100644
--- a/wrolpi/test/common.py
+++ b/wrolpi/test/common.py
@@ -10,7 +10,6 @@ from typing import List
import mock
import pytest
-from requests import Response
from wrolpi.common import set_test_media_directory, get_media_directory
from wrolpi.conftest import ROUTES_ATTACHED, test_db, test_client # noqa
@@ -128,7 +127,7 @@ class TestAPI(ExtendedTestCase):
def tearDown(self) -> None:
set_test_media_directory(None)
- def assertHTTPStatus(self, response: Response, status: int):
+ def assertHTTPStatus(self, response, status: int):
self.assertEqual(response.status_code, status)
assertOK = partialmethod(assertHTTPStatus, status=HTTPStatus.OK)
diff --git a/wrolpi/test/test_db.py b/wrolpi/test/test_db.py
index 8d2d43ca..821c9188 100644
--- a/wrolpi/test/test_db.py
+++ b/wrolpi/test/test_db.py
@@ -35,13 +35,3 @@ def test_optional_session(test_session):
func(session=test_session)
func(test_session)
-
-def test_test_session_is_required():
- """
- fixture `test_session` must be used in all tests to avoid using the real database.
- """
- try:
- with get_db_session() as session:
- raise Exception('test engine was not used!!')
- except ValueError as e:
- assert 'test database' in str(e)
diff --git a/wrolpi/test/test_downloader.py b/wrolpi/test/test_downloader.py
index 6e3423dd..8fdd773f 100644
--- a/wrolpi/test/test_downloader.py
+++ b/wrolpi/test/test_downloader.py
@@ -1,12 +1,14 @@
+import asyncio
from abc import ABC
from datetime import datetime
from itertools import zip_longest
+from typing import Tuple, Optional
from unittest import mock
from unittest.mock import MagicMock
import pytest
-from wrolpi.dates import local_timezone, Seconds
+from wrolpi.dates import local_timezone, Seconds, now
from wrolpi.db import get_db_context
from wrolpi.downloader import Downloader, Download, DownloadFrequency, DownloadResult
from wrolpi.errors import UnrecoverableDownloadError, InvalidDownload, WROLModeEnabled
@@ -35,6 +37,7 @@ class HTTPDownloader(Downloader):
return url.startswith('https://') or url.startswith('http://'), None
+@pytest.mark.asyncio
def test_register_downloader(test_session, test_download_manager):
"""Downloaders can be registered and have specific priorities."""
assert test_download_manager.instances == tuple()
@@ -54,84 +57,39 @@ def test_register_downloader(test_session, test_download_manager):
test_download_manager.register_downloader(permissive_downloader)
-def test_do_downloads(test_session, test_download_manager):
- _, session = get_db_context()
-
- http_downloader = HTTPDownloader()
- http_downloader.do_download = MagicMock()
- http_downloader.do_download.return_value = DownloadResult(success=True)
- test_download_manager.register_downloader(http_downloader)
-
- permissive_downloader = PermissiveDownloader(priority=100)
- permissive_downloader.do_download = MagicMock()
- permissive_downloader.do_download.return_value = DownloadResult(success=False)
- test_download_manager.register_downloader(permissive_downloader)
-
- # https is handled by the HTTP Downloader.
- test_download_manager.create_download('https://example.com', session)
- http_downloader.do_download.assert_called_once()
- permissive_downloader.do_download.assert_not_called()
- assert test_download_manager.get_download(session, 'https://example.com') is not None
-
- http_downloader.do_download.reset_mock()
-
- # try the permissive download, which returns a failure.
- test_download_manager.create_download('foo', session)
- http_downloader.do_download.assert_not_called()
- permissive_downloader.do_download.assert_called_once()
- download = test_download_manager.get_download(session, 'foo')
- assert download.attempts == 1
-
- # try again
- test_download_manager.create_download('foo', session)
- test_download_manager._do_downloads(session)
- download = test_download_manager.get_download(session, 'foo')
- assert download.attempts == 2
-
- # finally success
- permissive_downloader.do_download.return_value = DownloadResult(success=True)
- test_download_manager.create_download('foo', session)
- test_download_manager._do_downloads(session)
- download = test_download_manager.get_download(session, 'foo')
- assert download.status == 'complete'
-
- # No downloads left.
- assert list(test_download_manager.get_new_downloads(session)) == []
-
-
-def test_delete_old_once_downloads(test_session, test_download_manager):
+@pytest.mark.asyncio
+async def test_delete_old_once_downloads(test_session, test_download_manager):
"""Once-downloads over a month old should be deleted."""
with mock.patch('wrolpi.downloader.now') as mock_now:
mock_now.return_value = local_timezone(datetime(2020, 6, 5, 0, 0))
permissive_downloader = PermissiveDownloader(priority=0)
test_download_manager.register_downloader(permissive_downloader)
- _, session = get_db_context()
# Recurring downloads should not be deleted.
- d1 = test_download_manager.create_download('https://example.com/1', session, skip_download=True)
- d2 = test_download_manager.create_download('https://example.com/2', session, skip_download=True)
+ d1 = test_download_manager.create_download('https://example.com/1')
+ d2 = test_download_manager.create_download('https://example.com/2')
d1.frequency = 1
d2.frequency = 1
d2.started()
# Should be deleted.
- d3 = test_download_manager.create_download('https://example.com/3', session, skip_download=True)
- d4 = test_download_manager.create_download('https://example.com/4', session, skip_download=True)
+ d3 = test_download_manager.create_download('https://example.com/3')
+ d4 = test_download_manager.create_download('https://example.com/4')
d3.complete()
d4.complete()
d3.last_successful_download = local_timezone(datetime(2020, 1, 1, 0, 0, 0))
d4.last_successful_download = local_timezone(datetime(2020, 5, 1, 0, 0, 0))
# Not a month old.
- d5 = test_download_manager.create_download('https://example.com/5', session, skip_download=True)
+ d5 = test_download_manager.create_download('https://example.com/5')
d5.last_successful_download = local_timezone(datetime(2020, 6, 1, 0, 0, 0))
# An old, but pending download should not be deleted.
- d6 = test_download_manager.create_download('https://example.com/6', session, skip_download=True)
+ d6 = test_download_manager.create_download('https://example.com/6')
d6.last_successful_download = local_timezone(datetime(2020, 4, 1, 0, 0, 0))
d6.started()
test_download_manager.delete_old_once_downloads()
# Two old downloads are deleted.
- downloads = session.query(Download).order_by(Download.id).all()
+ downloads = test_session.query(Download).order_by(Download.id).all()
expected = [
dict(url='https://example.com/1', frequency=1, attempts=0, status='new'),
dict(url='https://example.com/2', frequency=1, attempts=1, status='pending'),
@@ -142,7 +100,8 @@ def test_delete_old_once_downloads(test_session, test_download_manager):
assert_dict_contains(download.dict(), expected)
-def test_ensure_download(test_session, test_download_manager):
+@pytest.mark.asyncio
+async def test_ensure_download(test_session, test_download_manager):
assert len(test_download_manager.get_downloads(test_session)) == 0
http_downloader = HTTPDownloader()
@@ -150,7 +109,7 @@ def test_ensure_download(test_session, test_download_manager):
http_downloader.do_download.return_value = DownloadResult(success=True)
test_download_manager.register_downloader(http_downloader)
- test_download_manager.create_download('https://example.com', test_session)
+ test_download_manager.create_download('https://example.com')
download = test_download_manager.get_download(test_session, 'https://example.com')
assert download.url == 'https://example.com'
@@ -159,19 +118,20 @@ def test_ensure_download(test_session, test_download_manager):
assert test_download_manager.get_download(test_session, 'https://example.com/not downloading') is None
-def test_create_downloads(test_session, test_download_manager):
+@pytest.mark.asyncio
+async def test_create_downloads(test_session, test_download_manager):
"""Multiple downloads can be scheduled using DownloadManager.create_downloads."""
http_downloader = HTTPDownloader(priority=0)
test_download_manager.register_downloader(http_downloader)
# Both URLs are valid and are scheduled.
- test_download_manager.create_downloads(['https://example.com/1', 'https://example.com/2'], skip_download=True)
+ test_download_manager.create_downloads(['https://example.com/1', 'https://example.com/2'])
downloads = test_download_manager.get_downloads(test_session)
assert {i.url for i in downloads} == {'https://example.com/1', 'https://example.com/2'}
# One URL is bad, neither should be scheduled.
with pytest.raises(InvalidDownload):
- test_download_manager.create_downloads(['https://example.com/3', 'bad url should fail'], skip_download=True)
+ test_download_manager.create_downloads(['https://example.com/3', 'bad url should fail'])
downloads = test_download_manager.get_downloads(test_session)
assert {i.url for i in downloads} == {'https://example.com/1', 'https://example.com/2'}
@@ -195,12 +155,13 @@ def test_downloader_must_have_name():
@pytest.mark.asyncio
async def test_download_wrol_mode(test_session, test_download_manager):
with pytest.raises(WROLModeEnabled):
- test_download_manager.do_downloads_sync()
+ await test_download_manager.do_downloads()
with pytest.raises(WROLModeEnabled):
- await test_download_manager.do_downloads_sync()
+ await test_download_manager.do_downloads()
-def test_download_get_downloader(test_session, test_download_manager, ):
+@pytest.mark.asyncio
+async def test_download_get_downloader(test_session, test_download_manager):
"""
A Download can find it's Downloader.
"""
@@ -209,14 +170,15 @@ def test_download_get_downloader(test_session, test_download_manager, ):
test_download_manager.register_downloader(permissive_downloader)
test_download_manager.register_downloader(http_downloader)
- download1 = test_download_manager.create_download('https://example.com', test_session, skip_download=True)
+ download1 = test_download_manager.create_download('https://example.com')
assert download1.get_downloader() == permissive_downloader
- download2 = test_download_manager.create_download('https://example.com', test_session, skip_download=True)
+ download2 = test_download_manager.create_download('https://example.com')
download2.downloader = 'bad downloader'
assert download2.get_downloader() is None
+@pytest.mark.asyncio
def test_calculate_next_download(test_session, test_download_manager, fake_now):
fake_now(datetime(2000, 1, 1))
download = Download()
@@ -249,79 +211,78 @@ def test_calculate_next_download(test_session, test_download_manager, fake_now):
assert test_download_manager.calculate_next_download(d3) == local_timezone(datetime(2000, 1, 9, 18))
-def test_recurring_downloads(test_session, test_download_manager, fake_now):
- """
- A recurring Download should be downloaded repeatedly forever.
- """
- _, session = get_db_context()
-
+@pytest.mark.asyncio
+async def test_recurring_downloads(test_session, test_download_manager, fake_now):
+ """A recurring Download should be downloaded forever."""
http_downloader = HTTPDownloader()
http_downloader.do_download = MagicMock()
http_downloader.do_download.return_value = DownloadResult(success=True)
test_download_manager.register_downloader(http_downloader)
# Download every hour.
- test_download_manager.recurring_download('https://example.com', Seconds.hour, skip_download=True)
+ test_download_manager.recurring_download('https://example.com', Seconds.hour)
# One download is scheduled.
- downloads = test_download_manager.get_new_downloads(session)
+ downloads = test_download_manager.get_new_downloads(test_session)
assert [(i.url, i.frequency) for i in downloads] == [('https://example.com', Seconds.hour)]
- now = fake_now(datetime(2020, 1, 1, 0, 0, 0))
+ now_ = fake_now(datetime(2020, 1, 1, 0, 0, 0))
# Download is processed and successful, no longer "new".
- test_download_manager.do_downloads_sync()
+ await test_download_manager.wait_for_all_downloads()
http_downloader.do_download.assert_called_once()
- assert list(test_download_manager.get_new_downloads(session)) == []
- downloads = list(test_download_manager.get_recurring_downloads(session))
+ assert list(test_download_manager.get_new_downloads(test_session)) == []
+ downloads = list(test_download_manager.get_recurring_downloads(test_session))
assert len(downloads) == 1
download = downloads[0]
expected = local_timezone(datetime(2020, 1, 1, 1, 0, 0))
assert download.next_download == expected
- assert download.last_successful_download == now
+ assert download.last_successful_download == now_
# Download is not due for an hour.
- test_download_manager.renew_recurring_downloads(session)
- assert list(test_download_manager.get_new_downloads(session)) == []
- assert download.last_successful_download == now
+ test_download_manager.renew_recurring_downloads(test_session)
+ await test_download_manager.wait_for_all_downloads()
+ assert list(test_download_manager.get_new_downloads(test_session)) == []
+ assert download.last_successful_download == now_
# Download is due an hour later.
fake_now(datetime(2020, 1, 1, 2, 0, 1))
- test_download_manager.renew_recurring_downloads(session)
- (download,) = list(test_download_manager.get_new_downloads(session))
+ test_download_manager.renew_recurring_downloads(test_session)
+ (download,) = list(test_download_manager.get_new_downloads(test_session))
# Download is "new" but has not been downloaded a second time.
assert download.next_download == expected
- assert download.last_successful_download == now
+ assert download.last_successful_download == now_
assert download.status == 'new'
# Try the download, but it fails.
http_downloader.do_download.reset_mock()
http_downloader.do_download.return_value = DownloadResult(success=False)
- test_download_manager.do_downloads_sync()
+ await test_download_manager.wait_for_all_downloads()
http_downloader.do_download.assert_called_once()
- download = session.query(Download).one()
+ download = test_session.query(Download).one()
# Download is deferred, last successful download remains the same.
assert download.status == 'deferred'
- assert download.last_successful_download == now
+ assert download.last_successful_download == now_
# Download should be retried after the DEFAULT_RETRY_FREQUENCY.
expected = local_timezone(datetime(2020, 1, 1, 3, 0, 0, 997200))
assert download.next_download == expected
# Try the download again, it finally succeeds.
http_downloader.do_download.reset_mock()
- now = fake_now(datetime(2020, 1, 1, 4, 0, 1))
+ now_ = fake_now(datetime(2020, 1, 1, 4, 0, 1))
http_downloader.do_download.return_value = DownloadResult(success=True)
- test_download_manager.renew_recurring_downloads(session)
- test_download_manager.do_downloads_sync()
+ test_download_manager.renew_recurring_downloads(test_session)
+ await test_download_manager.wait_for_all_downloads()
http_downloader.do_download.assert_called_once()
- download = session.query(Download).one()
+ download = test_session.query(Download).one()
assert download.status == 'complete'
- assert download.last_successful_download == now
+ assert download.last_successful_download == now_
# Floats cause slightly wrong date.
assert download.next_download == local_timezone(datetime(2020, 1, 1, 5, 0, 0, 997200))
-def test_max_attempts(test_session, test_download_manager):
+@pytest.mark.asyncio
+async def test_max_attempts(test_session, test_download_manager):
"""A Download will only be attempted so many times, it will eventually be deleted."""
_, session = get_db_context()
@@ -330,23 +291,27 @@ def test_max_attempts(test_session, test_download_manager):
http_downloader.do_download.return_value = DownloadResult(success=True)
test_download_manager.register_downloader(http_downloader)
- test_download_manager.create_download('https://example.com', session)
+ test_download_manager.create_download('https://example.com')
+ await test_download_manager.wait_for_all_downloads()
download = session.query(Download).one()
assert download.attempts == 1
- test_download_manager.create_download('https://example.com', session)
+ test_download_manager.create_download('https://example.com')
+ await test_download_manager.wait_for_all_downloads()
download = session.query(Download).one()
assert download.attempts == 2
# There are no further attempts.
http_downloader.do_download.side_effect = UnrecoverableDownloadError()
- test_download_manager.create_download('https://example.com', session)
+ test_download_manager.create_download('https://example.com')
+ await test_download_manager.wait_for_all_downloads()
download = session.query(Download).one()
assert download.attempts == 3
assert download.status == 'failed'
-def test_skip_urls(test_session, test_download_manager, test_client, assert_download_urls):
+@pytest.mark.asyncio
+async def test_skip_urls(test_session, test_download_manager, test_client, assert_download_urls):
"""The DownloadManager will not create downloads for URLs in it's skip list."""
_, session = get_db_context()
from wrolpi.downloader import DOWNLOAD_MANAGER_CONFIG
@@ -361,41 +326,79 @@ def test_skip_urls(test_session, test_download_manager, test_client, assert_down
'https://example.com/1',
'https://example.com/skipme',
'https://example.com/2',
- ], downloader_name=HTTPDownloader.name)
+ ], downloader=HTTPDownloader.name)
+ await test_download_manager.wait_for_all_downloads()
assert_download_urls({'https://example.com/1', 'https://example.com/2'})
assert DOWNLOAD_MANAGER_CONFIG.skip_urls == ['https://example.com/skipme']
test_download_manager.delete_completed()
# The user can start a download even if the URL is in the skip list.
- test_download_manager.create_download('https://example.com/skipme', session=test_session, reset_attempts=True)
+ test_download_manager.create_download('https://example.com/skipme', reset_attempts=True)
assert_download_urls({'https://example.com/skipme'})
assert DOWNLOAD_MANAGER_CONFIG.skip_urls == []
-def test_process_runner_timeout(test_directory):
+@pytest.mark.asyncio
+async def test_process_runner_timeout(test_directory):
"""A Downloader can cancel it's download using a timeout."""
# Default timeout of 3 seconds.
downloader = Downloader(0, 'downloader', timeout=3)
# Default timeout is obeyed.
start = datetime.now()
- downloader.process_runner('https://example.com', ('sleep', '8'), test_directory)
+ await downloader.process_runner('https://example.com', ('sleep', '8'), test_directory)
elapsed = datetime.now() - start
assert 3 < elapsed.total_seconds() < 4
# One-off timeout is obeyed.
start = datetime.now()
- downloader.process_runner('https://example.com', ('sleep', '8'), test_directory, timeout=1)
+ await downloader.process_runner('https://example.com', ('sleep', '8'), test_directory, timeout=1)
elapsed = datetime.now() - start
assert 1 < elapsed.total_seconds() < 2
- # This Downloader's timeout is ignored in favor of the global timeout.
+ # Global timeout is obeyed.
from wrolpi.common import WROLPI_CONFIG
WROLPI_CONFIG.download_timeout = 3
-
- # Global timeout is obeyed.
start = datetime.now()
- downloader.process_runner('https://example.com', ('sleep', '8'), test_directory)
+ await downloader.process_runner('https://example.com', ('sleep', '8'), test_directory)
elapsed = datetime.now() - start
assert 2 < elapsed.total_seconds() < 4
+
+
+@pytest.mark.asyncio
+async def test_multi_domain_download(test_session, test_directory, test_download_manager):
+ class TestDownloader(Downloader, ABC):
+ name = 'test_downloader'
+
+ @classmethod
+ def valid_url(cls, url: str) -> Tuple[bool, Optional[dict]]:
+ return True, {}
+
+ async def do_download(self, download: Download) -> DownloadResult:
+ await asyncio.sleep(1)
+ return DownloadResult(success=True)
+
+ async def wait_and_assert(urls_, max_seconds):
+ downloads = test_download_manager.create_downloads(urls_, downloader=test_downloader.name)
+ assert downloads and all(isinstance(i, Download) for i in downloads), 'Did not get Downloads'
+ start = now()
+ await test_download_manager.wait_for_all_downloads()
+ assert (elapsed := now() - start).total_seconds() < max_seconds, \
+ f'Downloads took {elapsed} seconds. Were they in parallel?'
+ downloads = test_download_manager.get_downloads(test_session)
+ incomplete = [i.url for i in downloads if i.status != 'complete']
+ assert not incomplete, f'Some downloads did not complete: {incomplete}'
+
+ test_downloader = TestDownloader()
+ test_download_manager.register_downloader(test_downloader)
+
+ # Each domain's downloads will be processed one at a time.
+ # example.com and example.org will both download in parallel.
+ urls = ['https://example.com/1', 'https://example.com/2', 'https://example.com/3', 'https://example.org/1',
+ 'https://example.org/2']
+ await wait_and_assert(urls, 4)
+
+ # Downloading more domains that workers is possible.
+ urls = [f'https://example.{i}' for i in range(test_download_manager.worker_count + 2)]
+ await wait_and_assert(urls, 3)
diff --git a/wrolpi/test/test_root_api.py b/wrolpi/test/test_root_api.py
index df488018..120398b3 100644
--- a/wrolpi/test/test_root_api.py
+++ b/wrolpi/test/test_root_api.py
@@ -15,6 +15,7 @@ from wrolpi.test.common import TestAPI, wrap_test_db, skip_circleci
class TestRootAPI(TestAPI):
+ @wrap_test_db
def test_index(self):
"""
Index should have some details in an HTML response
@@ -61,6 +62,7 @@ class TestRootAPI(TestAPI):
assert response.status_code == HTTPStatus.BAD_REQUEST
assert json.loads(response.body) == {'valid': False, 'regex': '.*(missing parenthesis.*'}
+ @wrap_test_db
def test_get_settings(self):
with mock.patch('wrolpi.admin.CPUFREQ_INFO_BIN', '/usr/bin/cpufreq-info'), \
mock.patch('wrolpi.admin.NMCLI_BIN', '/usr/bin/nmcli'):
diff --git a/wrolpi/test/test_rss.py b/wrolpi/test/test_rss.py
index 962cf337..df984b7f 100644
--- a/wrolpi/test/test_rss.py
+++ b/wrolpi/test/test_rss.py
@@ -4,6 +4,7 @@ from itertools import zip_longest
from typing import List
import mock
+import pytest
from wrolpi.db import optional_session
from wrolpi.downloader import Download, DownloadResult, RSSDownloader
@@ -20,13 +21,8 @@ class RSSHTTPDownloader(HTTPDownloader):
return download
-def test_rss_download_invalid(test_session, test_download_manager):
- """Feeds are not downloaded when FeedParser reports an issue."""
- rss_downloader = RSSDownloader()
- test_download_manager.register_downloader(rss_downloader)
-
-
-def test_rss_download(test_session, test_download_manager):
+@pytest.mark.asyncio
+async def test_rss_download(test_session, test_download_manager):
"""An RSS Downloader will create new Downloads for every link in the feed."""
rss_downloader = RSSDownloader()
test_download_manager.register_downloader(rss_downloader)
@@ -50,7 +46,8 @@ def test_rss_download(test_session, test_download_manager):
dict(link='https://example.com/c'),
]
)
- test_download_manager.create_download('https://example.com/feed', test_session, sub_downloader='http')
+ test_download_manager.create_download('https://example.com/feed', sub_downloader='http')
+ await test_download_manager.wait_for_all_downloads()
# Feed download is complete.
check_downloads([
@@ -68,7 +65,8 @@ def test_rss_download(test_session, test_download_manager):
dict(link='https://example.com/d'),
]
)
- test_download_manager.create_download('https://example.com/feed', test_session, sub_downloader='http')
+ test_download_manager.create_download('https://example.com/feed', sub_downloader='http')
+ await test_download_manager.wait_for_all_downloads()
# Only the new URLs are Archived.
check_downloads([
@@ -80,28 +78,8 @@ def test_rss_download(test_session, test_download_manager):
])
-def test_post_rss_download(test_session, test_client, test_download_manager):
- """An RSS can be downloaded in the UI"""
- rss_downloader = RSSDownloader()
- test_download_manager.register_downloader(rss_downloader)
- http_downloader = RSSHTTPDownloader()
- test_download_manager.register_downloader(http_downloader)
-
- with mock.patch('wrolpi.downloader.parse_feed') as mock_parse_feed:
- mock_parse_feed.return_value = dict(bozo=0, entries=[])
- content = dict(urls='https://example.com/feed', downloader='rss', frequency=100)
- request, response = test_client.post('/api/download', content=json.dumps(content))
- assert response.status == HTTPStatus.NO_CONTENT
-
- # Download was attempted.
- (download,) = test_download_manager.get_downloads(test_session)
- assert download.url == 'https://example.com/feed'
- assert download.frequency == 100
- assert download.downloader == 'rss'
- assert download.attempts == 1
-
-
-def test_rss_no_entries(test_session, test_download_manager):
+@pytest.mark.asyncio
+async def test_rss_no_entries(test_session, test_download_manager):
"""An RSS feed with no entries is handled."""
rss_downloader = RSSDownloader()
test_download_manager.register_downloader(rss_downloader)
@@ -109,6 +87,7 @@ def test_rss_no_entries(test_session, test_download_manager):
with mock.patch('wrolpi.downloader.parse_feed') as mock_parse_feed:
mock_parse_feed.return_value = dict(bozo=0, ) # missing `entries`
test_download_manager.create_download('https://example.com/feed', test_session, sub_downloader='http')
+ await test_download_manager.wait_for_all_downloads()
(download,) = test_download_manager.get_downloads(test_session)
assert 'entries' in download.error
diff --git a/wrolpi/test/test_timeout.py b/wrolpi/test/test_timeout.py
deleted file mode 100644
index a84d098c..00000000
--- a/wrolpi/test/test_timeout.py
+++ /dev/null
@@ -1,47 +0,0 @@
-import time
-import unittest
-from unittest import mock
-
-from wrolpi.timeout import timeout, WorkerException
-
-
-class TestTimeout(unittest.TestCase):
-
- def test_timeout(self):
- """
- The process will be killed if it takes too long.
- """
-
- @timeout(2)
- def func():
- time.sleep(5)
- raise Exception('Was not killed!')
-
- self.assertRaises(TimeoutError, func)
-
- def test_timeout_exception(self):
- """
- If the wrapped function has an error, it will raise a WorkerException.
- """
-
- @timeout(10)
- def func():
- raise ValueError('Cause a worker error!')
-
- self.assertRaises(WorkerException, func)
-
- def test_test_timeout(self):
- """Tests can overwrite the timeout"""
-
- @timeout(0.1)
- def func():
- time.sleep(1)
- return True
-
- # Killed after 0.1 seconds
- self.assertRaises(TimeoutError, func)
-
- with mock.patch('wrolpi.timeout.TEST_TIMEOUT', 2):
- # New timeout allows the function to finish.
- result = func()
- self.assertTrue(result)
diff --git a/wrolpi/timeout.py b/wrolpi/timeout.py
deleted file mode 100644
index 3322014e..00000000
--- a/wrolpi/timeout.py
+++ /dev/null
@@ -1,114 +0,0 @@
-"""
-Wrap a function with `timeout` to kill it after some amount of seconds, the wrapped function is run in it's own
-process. If the timeout is reached, kill the process and raise a TimeoutError.
-
-Log any errors in the global logger. If the worker process raises an Exception, raise an Exception sharing it's
-description, log the traceback.
-
->>> @timeout(5)
->>> def long_running_function():
->>> while True:
->>> # Run infinitely
->>> time.sleep(1)
-
->>> long_running_function()
-TimeoutError()
-"""
-import multiprocessing.connection
-import time
-import traceback
-from datetime import timedelta
-from functools import wraps
-from multiprocessing import Process, Pipe
-
-from wrolpi.common import logger
-from wrolpi.dates import utc_now
-
-logger = logger.getChild(__name__)
-
-
-class WorkerException(Exception):
- pass
-
-
-TEST_TIMEOUT = None
-KILL_ATTEMPTS = 100
-
-
-def timeout(seconds: float) -> callable:
- """
- Create a worker process, let it run until the timeout, then kill it. If the worker isn't killed, return
- it's result or exception. If it is killed, raise a TimeoutError.
- """
-
- def wrapper(func):
-
- def worker_wrapper(pipe: multiprocessing.connection.Connection, *args, **kwargs):
- """
- Wraps a worker function so we can catch it's Exception, or return its results.
- """
- to_send = {}
- try:
- result = func(*args, **kwargs)
- to_send['result'] = result
- except Exception as e:
- to_send.update({'exception': str(e), 'traceback': traceback.format_exc()})
- finally:
- pipe.send(to_send)
-
- @wraps(func)
- def wrapped(*args, **kwargs):
- # unittest may override the timeout
- _seconds = TEST_TIMEOUT or seconds
-
- # Start the worker_wrapper, give it a way to communicate back.
- parent_conn, child_conn = Pipe()
- args = (child_conn,) + args
- worker = Process(target=worker_wrapper, args=args, kwargs=kwargs)
- worker.start()
-
- # Continually attempt to get a result until the timeout is reached.
- kill_time = utc_now() + timedelta(seconds=_seconds)
- while True:
- if not worker.is_alive():
- # Worker exited.
- break
- if not utc_now() < kill_time:
- # Timeout was reached.
- break
- if parent_conn.poll(0.1): # This prevents a "busy loop"
- # Result is waiting.
- break
-
- # Get the result, if any.
- result = None
- if parent_conn.poll(0):
- result = parent_conn.recv()
-
- attempts = 0
- while worker.is_alive():
- # Worker has run out of time.
- worker.kill()
- time.sleep(0.1)
- attempts += 1
- if attempts > KILL_ATTEMPTS:
- logger.warning(f'Failed to kill worker! {func}')
- break
-
- worker.join()
-
- if result:
- if 'result' in result:
- return result['result']
- elif 'exception' in result:
- exception = result['exception']
- _traceback = result['traceback']
- logger.warning(f'Traceback from worker: {_traceback}')
- raise WorkerException(f'Exception from worker: {exception}')
-
- if attempts > 0:
- raise TimeoutError(f'Killed worker after {_seconds} seconds: {func}')
-
- return wrapped
-
- return wrapper
diff --git a/wrolpi/version.py b/wrolpi/version.py
index dabeb96b..11320bab 100644
--- a/wrolpi/version.py
+++ b/wrolpi/version.py
@@ -1,4 +1,4 @@
-__version__ = '0.5.3-beta'
+__version__ = '0.6.0-beta'
import subprocess