diff --git a/modules/archive/test/test_lib.py b/modules/archive/test/test_lib.py index 39fe6da6..e868be84 100644 --- a/modules/archive/test/test_lib.py +++ b/modules/archive/test/test_lib.py @@ -51,35 +51,6 @@ class TestArchive(TestAPI): assert str(j).endswith('archive/example.com/2001-01-01 00:00:00.000000-readability.json') assert str(c).endswith('archive/example.com/2001-01-01 00:00:00.000000.png') - @wrap_test_db - def test_new_archive(self): - with mock.patch('modules.archive.lib.request_archive', make_fake_request_archive()): - archive1 = do_archive('https://example.com') - # Everything is filled out. - self.assertIsInstance(archive1, Archive) - self.assertIsNotNone(archive1.archive_datetime) - self.assertIsInstance(archive1.singlefile_path, MediaPath) - self.assertIsInstance(archive1.readability_path, MediaPath) - self.assertIsInstance(archive1.readability_txt_path, MediaPath) - self.assertIsInstance(archive1.screenshot_path, MediaPath) - self.assertEqual(archive1.title, 'ジにてこちら') - self.assertIsNotNone(archive1.url) - self.assertIsNotNone(archive1.domain) - - # The actual files were dumped and read correctly. - with open(archive1.singlefile_path.path) as fh: - self.assertEqual(fh.read(), '\ntest single-file\nジにてこちら\nsome title') - with open(archive1.readability_path.path) as fh: - self.assertEqual(fh.read(), 'test readability content') - with open(archive1.readability_txt_path.path) as fh: - self.assertEqual(fh.read(), 'test readability textContent') - with open(archive1.readability_json_path.path) as fh: - self.assertEqual(json.load(fh), {'title': 'ジにてこちら', 'url': 'https://example.com'}) - - archive2 = do_archive('https://example.com') - # Domain is reused. - self.assertEqual(archive1.domain, archive2.domain) - @wrap_test_db def test_no_screenshot(self): with mock.patch('modules.archive.lib.request_archive', make_fake_request_archive(screenshot=False)): @@ -132,39 +103,6 @@ class TestArchive(TestAPI): except ValueError as e: self.assertIn('relative', str(e), f'Relative path error was not raised') - @wrap_test_db - def test_get_title_from_html(self): - with wrap_media_directory(): - with mock.patch('modules.archive.lib.request_archive', make_fake_request_archive()): - archive = do_archive('example.com') - self.assertEqual(archive.title, 'ジにてこちら') - - def fake_request_archive(_): - singlefile = '\ntest single-file\nジにてこちら\nsome title' - r = dict( - content=f'test readability content', - textContent='test readability textContent', - ) - s = b'screenshot data' - return singlefile, r, s - - with mock.patch('modules.archive.lib.request_archive', fake_request_archive): - archive = do_archive('example.com') - self.assertEqual(archive.title, 'some title') - - def fake_request_archive(_): - singlefile = '' - r = dict( - content=f'missing a title', - textContent='', - ) - s = b'screenshot data' - return singlefile, r, s - - with mock.patch('modules.archive.lib.request_archive', fake_request_archive): - archive = do_archive('example.com') - self.assertIsNone(archive.title) - @wrap_test_db def test_refresh_archives(self): with wrap_media_directory(): @@ -324,3 +262,69 @@ def test_get_domains(test_session, archive_factory): archive3.delete() test_session.commit() assert [i['domain'] for i in get_domains()] == [] + + +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') + # Everything is filled out. + assert isinstance(archive1, Archive) + assert archive1.archive_datetime is not None + assert isinstance(archive1.singlefile_path, MediaPath) + assert isinstance(archive1.readability_path, MediaPath) + assert isinstance(archive1.readability_txt_path, MediaPath) + assert isinstance(archive1.screenshot_path, MediaPath) + assert archive1.title == 'ジにてこちら' + assert archive1.url is not None + assert archive1.domain is not None + + # The actual files were dumped and read correctly. + with open(archive1.singlefile_path.path) as fh: + assert fh.read() == '\ntest single-file\nジにてこちら\nsome title' + with open(archive1.readability_path.path) as fh: + assert fh.read() == 'test readability content' + with open(archive1.readability_txt_path.path) as fh: + assert fh.read() == 'test readability textContent' + with open(archive1.readability_json_path.path) as fh: + assert json.load(fh) == {'title': 'ジにてこちら', 'url': 'https://example.com'} + + fake_now(datetime(2000, 1, 2)) + archive2 = do_archive('https://example.com') + # Domain is reused. + assert archive1.domain == archive2.domain + + +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') + assert archive.title == 'ジにてこちら' + + def fake_request_archive(_): + singlefile = '\ntest single-file\nジにてこちら\nsome title' + r = dict( + content=f'test readability content', + textContent='test readability textContent', + ) + s = b'screenshot data' + return singlefile, r, s + + fake_now(datetime(2000, 1, 2)) + with mock.patch('modules.archive.lib.request_archive', fake_request_archive): + archive = do_archive('example.com') + assert archive.title == 'some title' + + def fake_request_archive(_): + singlefile = '' + r = dict( + content=f'missing a title', + textContent='', + ) + s = b'screenshot data' + return singlefile, r, s + + fake_now(datetime(2000, 1, 3)) + with mock.patch('modules.archive.lib.request_archive', fake_request_archive): + archive = do_archive('example.com') + assert archive.title is None diff --git a/modules/videos/__init__.py b/modules/videos/__init__.py index 5cd95bbb..8bfe14d6 100644 --- a/modules/videos/__init__.py +++ b/modules/videos/__init__.py @@ -1,8 +1 @@ from modules.videos import downloader -from wrolpi import before_startup - - -@before_startup -def startup_spread_channel_downloads(): - from .channel import lib as channel_lib - channel_lib.spread_channel_downloads() diff --git a/modules/videos/channel/lib.py b/modules/videos/channel/lib.py index 64e329ec..2bf3059c 100644 --- a/modules/videos/channel/lib.py +++ b/modules/videos/channel/lib.py @@ -2,6 +2,7 @@ from collections import defaultdict from datetime import timedelta from typing import List, Dict, Union, Generator +from sqlalchemy.orm import Session from sqlalchemy.orm.exc import NoResultFound from wrolpi.common import sanitize_link, run_after, get_relative_to_media_directory, make_media_directory, logger @@ -115,33 +116,12 @@ def _spread_by_frequency(channels: List[Channel]) -> Generator[Dict, None, None] yield dict(url=channel.url, frequency=frequency, next_download=next_download) -def spread_channel_downloads(): - """ - Channels should be downloaded in a manner that is spread out over their frequency. For example, three channels - with a frequency of a week should be downloaded on different days that week. - """ - with get_db_session(commit=True) as session: - channels = list(session.query(Channel).filter( - Channel.url != None, - Channel.url != '', - Channel.download_frequency != None - ).order_by(Channel.link).all()) # noqa - - url_next_download = _spread_by_frequency(channels) - - for info in url_next_download: - url, frequency, next_download = info['url'], info['frequency'], info['next_download'] - download = download_manager.get_or_create_download(url, session) - download.downloader = 'video_channel' - download.frequency = frequency - download.next_download = next_download - - @run_after(save_channels_config) -def update_channel(data, link): +def update_channel(data: dict, link: str) -> Channel: + """Update a Channel's DB record""" with get_db_session(commit=True) as session: try: - channel = session.query(Channel).filter_by(link=link).one() + channel: Channel = session.query(Channel).filter_by(link=link).one() except NoResultFound: raise UnknownChannel() @@ -184,13 +164,10 @@ def update_channel(data, link): # Apply the changes now that we've OK'd them channel.update(data) - spread_channel_downloads() - return channel @run_after(save_channels_config) -@run_after(spread_channel_downloads) def create_channel(data: dict, return_dict: bool = True) -> Union[Channel, dict]: """ Create a new Channel. Check for conflicts with existing Channels. @@ -242,12 +219,12 @@ def delete_channel(link): def download_channel(link: str): - """ - Create a Download record for a Channel's entire catalog. - """ + """Create a Download record for a Channel's entire catalog. Start downloading.""" channel = get_channel(link=link, return_dict=False) - with get_db_session(commit=True) as session: - if not download_manager.get_download(session, channel.url): - raise InvalidDownload(f'Channel {channel.name} does not have a download!') - - download_manager.create_download(channel.url, session, downloader='video_channel') + session = Session.object_session(channel) + download = channel.get_download() + if not download: + 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 a610bb92..6e30d309 100644 --- a/modules/videos/channel/test/test_api.py +++ b/modules/videos/channel/test/test_api.py @@ -10,7 +10,7 @@ import mock import pytest from modules.videos.api import refresh_queue -from modules.videos.channel.lib import spread_channel_downloads, delete_channel +from modules.videos.channel.lib import delete_channel from modules.videos.common import get_no_channel_directory from modules.videos.lib import upsert_video from modules.videos.models import Channel, Video @@ -380,7 +380,6 @@ class TestVideoAPI(TestAPI): location = response.headers['Location'] # No downloads are scheduled. - spread_channel_downloads() with get_db_session() as session: downloads = list(download_manager.get_new_downloads(session)) self.assertEqual(len(downloads), 0) diff --git a/modules/videos/conftest.py b/modules/videos/conftest.py index c4762c1c..b25edee4 100644 --- a/modules/videos/conftest.py +++ b/modules/videos/conftest.py @@ -7,11 +7,10 @@ import mock import pytest from PIL import Image -from modules.videos.channel.lib import spread_channel_downloads from modules.videos.downloader import VideoDownloader, ChannelDownloader from modules.videos.models import Channel, Video from wrolpi.common import sanitize_link -from wrolpi.downloader import DownloadFrequency, DownloadManager +from wrolpi.downloader import DownloadFrequency, DownloadManager, Download from wrolpi.vars import PROJECT_DIR @@ -57,7 +56,7 @@ def channel_factory(test_session, test_directory): @pytest.fixture -def download_channel(test_session, test_directory) -> Channel: +def download_channel(test_session, test_directory, video_download_manager) -> Channel: """ Get a test Channel that has a download frequency. """ @@ -69,9 +68,9 @@ def download_channel(test_session, test_directory) -> Channel: download_frequency=DownloadFrequency.weekly, link='downloadchannel', ) - test_session.add(channel) + download = Download(url=channel.url, downloader='video_channel', frequency=channel.download_frequency) + test_session.add_all([channel, download]) test_session.commit() - spread_channel_downloads() return channel diff --git a/modules/videos/downloader.py b/modules/videos/downloader.py index 8b8138d8..0e072ef2 100755 --- a/modules/videos/downloader.py +++ b/modules/videos/downloader.py @@ -93,7 +93,7 @@ class ChannelDownloader(Downloader, ABC): YT_DLP_BIN = pathlib.Path('/usr/local/bin/yt-dlp') # Location in docker container if not YT_DLP_BIN.is_file(): YT_DLP_BIN = pathlib.Path('/opt/wrolpi/venv/bin/yt-dlp') # Use virtual environment location -if not YT_DLP_BIN.is_file(): +if not YT_DLP_BIN.is_file() and not PYTEST: logger.error('COULD NOT FIND YT-DLP!!!') diff --git a/modules/videos/lib.py b/modules/videos/lib.py index 053bf500..499bd484 100644 --- a/modules/videos/lib.py +++ b/modules/videos/lib.py @@ -391,7 +391,7 @@ async def get_statistics(): return ret -NAME_PARSER = re.compile(r'(.*?)_((?:\d+?)|(?:NA))_(?:(.+?)_)?(.*)\.' +NAME_PARSER = re.compile(r'(.*?)_((?:\d+?)|(?:NA))_(?:(.{5,15})_)?(.*)\.' r'(jpg|webp|flv|mp4|part|info\.json|description|webm|..\.srt|..\.vtt)') diff --git a/modules/videos/models.py b/modules/videos/models.py index b9ca0190..36a597e5 100644 --- a/modules/videos/models.py +++ b/modules/videos/models.py @@ -11,8 +11,8 @@ from sqlalchemy.orm.collections import InstrumentedList from wrolpi.common import Base, tsvector, ModelHelper, logger, get_media_directory from wrolpi.dates import now, TZDateTime from wrolpi.db import get_db_curs -from wrolpi.downloader import Download -from wrolpi.errors import UnknownVideo, UnknownFile, UnknownDirectory +from wrolpi.downloader import Download, download_manager +from wrolpi.errors import UnknownVideo, UnknownFile, UnknownDirectory, InvalidDownload from wrolpi.media_path import MediaPathType, MediaPath logger = logger.getChild(__name__) @@ -302,9 +302,20 @@ class Channel(ModelHelper, Base): session.delete(self) def update(self, data: dict): + """ + Update the values of this Channel. Will also update the Channel's Download, if it has one. + """ for key, value in data.items(): setattr(self, key, value) + if not self.url: + return + download = self.get_download() + if download: + session = Session.object_session(self) + download.frequency = self.download_frequency + download.next_download = download_manager.get_next_download(download, session) + def config_view(self) -> dict: """ Retrieve the data about this Channel that should be stored in a config file. @@ -334,7 +345,7 @@ class Channel(ModelHelper, Base): Get the Download row for this Channel. If there isn't a Download, return None. """ if not self.url: - raise ValueError(f'Channel {self.name} does not have a URL to download!') + raise InvalidDownload(f'Channel {self.name} does not have a URL to download!') session = Session.object_session(self) download = session.query(Download).filter_by(url=self.url).one_or_none() diff --git a/modules/videos/test/test_downloader.py b/modules/videos/test/test_downloader.py index 7aa2d90b..60684bde 100644 --- a/modules/videos/test/test_downloader.py +++ b/modules/videos/test/test_downloader.py @@ -1,19 +1,17 @@ import pathlib from copy import copy -from datetime import datetime from unittest import mock import pytest from yt_dlp.utils import UnsupportedError -from modules.videos.channel.lib import spread_channel_downloads, download_channel +from modules.videos.channel.lib import download_channel from modules.videos.downloader import find_all_missing_videos, VideoDownloader, \ ChannelDownloader, get_or_create_channel from modules.videos.models import Channel, Video from modules.videos.test.common import create_channel_structure -from wrolpi.dates import local_timezone from wrolpi.db import get_db_session, get_db_context -from wrolpi.downloader import DownloadManager, Download, DownloadFrequency +from wrolpi.downloader import DownloadManager, Download from wrolpi.errors import InvalidDownload from wrolpi.test.common import wrap_test_db, TestAPI from wrolpi.test.test_downloader import HTTPDownloader @@ -286,53 +284,6 @@ class TestVideosDownloaders(TestAPI): download = self.mgr.get_downloads(session)[0] assert download.status == 'failed' - @wrap_test_db - @mock.patch('modules.videos.channel.lib.today', lambda: local_timezone(datetime(2020, 1, 1, 0, 0, 0))) - def test_spread_channel_downloads(self): - with get_db_session(commit=True) as session: - channel1 = Channel(name='channel1', link='channel1', url='https://example.com/1', - download_frequency=DownloadFrequency.weekly) - channel2 = Channel(name='channel2', link='channel2', url='https://example.com/2', - download_frequency=DownloadFrequency.weekly) - session.add_all([channel1, channel2]) - - def check_frequencies(expected): - with get_db_session(commit=True) as session: - downloads = list(session.query(Download).all()) - self.assertLength(downloads, expected) - for download in downloads: - next_download = expected[download.url] - self.assertEqual(next_download, str(download.next_download), - f'{download.url} frequency {next_download} != {download.next_download}') - - # Spread the downloads out. - spread_channel_downloads() - - expected = { - 'https://example.com/1': '2020-01-08 00:00:00-07:00', - 'https://example.com/2': '2020-01-11 12:00:00-07:00', - } - check_frequencies(expected) - - with get_db_session(commit=True) as session: - channel3 = Channel(name='channel3', link='channel3', url='https://example.com/3', - download_frequency=DownloadFrequency.daily) - channel4 = Channel(name='channel4', link='channel4', url='https://example.com/4', - download_frequency=DownloadFrequency.weekly) - # 5 and 6 should be ignored. - channel5 = Channel(name='channel5', link='channel5', url='https://example.com/5') - channel6 = Channel(name='channel6', link='channel6', download_frequency=DownloadFrequency.daily) - session.add_all([channel3, channel4, channel5, channel6]) - - spread_channel_downloads() - expected = { - 'https://example.com/1': '2020-01-08 00:00:00-07:00', - 'https://example.com/2': '2020-01-10 08:00:00-07:00', - 'https://example.com/3': '2020-01-02 00:00:00-07:00', - 'https://example.com/4': '2020-01-12 16:00:00-07:00', - } - check_frequencies(expected) - @wrap_test_db def test_get_or_create_channel(self): """ diff --git a/modules/videos/test/test_lib.py b/modules/videos/test/test_lib.py index b16a66f5..694382bc 100644 --- a/modules/videos/test/test_lib.py +++ b/modules/videos/test/test_lib.py @@ -7,80 +7,13 @@ from unittest import mock import pytest import sqlalchemy -from modules.videos.channel import lib from modules.videos.common import apply_info_json from modules.videos.lib import validate_videos, parse_video_file_name, upsert_video -from modules.videos.models import Channel, Video +from modules.videos.models import Video from modules.videos.video.lib import video_search from wrolpi.dates import local_timezone -from wrolpi.downloader import DownloadFrequency from wrolpi.vars import PROJECT_DIR -daily, weekly, days30 = DownloadFrequency.daily, DownloadFrequency.weekly, DownloadFrequency.days30 - - -@pytest.mark.parametrize( - 'channels,expected', - [ - ([Channel(download_frequency=daily, url='https://example.com/1')], - [('https://example.com/1', daily, '2000-01-02 00:00:00-07:00')] - ), - # Several different download frequencies. - ([ - Channel(download_frequency=daily, url='https://example.com/1'), - Channel(download_frequency=daily, url='https://example.com/2'), - Channel(download_frequency=daily, url='https://example.com/3'), - Channel(download_frequency=weekly, url='https://example.com/4'), - Channel(download_frequency=weekly, url='https://example.com/5'), - Channel(download_frequency=days30, url='https://example.com/6'), - ], - [ - ('https://example.com/1', daily, '2000-01-02 00:00:00-07:00'), - ('https://example.com/2', daily, '2000-01-02 08:00:00-07:00'), - ('https://example.com/3', daily, '2000-01-02 16:00:00-07:00'), - ('https://example.com/4', weekly, '2000-01-08 00:00:00-07:00'), - ('https://example.com/5', weekly, '2000-01-11 12:00:00-07:00'), - ('https://example.com/6', days30, '2000-01-31 00:00:00-07:00'), - ] - ), - # Multiple downloads in a day. - ([Channel(download_frequency=daily, url=f'https://example.com/{i}') for i in range(1, 3)], - [ - ('https://example.com/1', daily, '2000-01-02 00:00:00-07:00'), - ('https://example.com/2', daily, '2000-01-02 12:00:00-07:00'), - ] - ), - # More downloads in a week than days in a week. - ([Channel(download_frequency=weekly, url=f'https://example.com/{i}') for i in range(1, 9)], - [ - ('https://example.com/1', weekly, '2000-01-08 00:00:00-07:00'), - ('https://example.com/2', weekly, '2000-01-08 21:00:00-07:00'), - ('https://example.com/3', weekly, '2000-01-09 18:00:00-07:00'), - ('https://example.com/4', weekly, '2000-01-10 15:00:00-07:00'), - ('https://example.com/5', weekly, '2000-01-11 12:00:00-07:00'), - ('https://example.com/6', weekly, '2000-01-12 09:00:00-07:00'), - ('https://example.com/7', weekly, '2000-01-13 06:00:00-07:00'), - ('https://example.com/8', weekly, '2000-01-14 03:00:00-07:00'), - ] - ), - # Any Channel without a URL AND frequency is ignored. - ([ - Channel(download_frequency=daily, url='https://example.com/1'), - Channel(download_frequency=daily), - Channel(url='this should be ignored'), - ], - [('https://example.com/1', daily, '2000-01-02 00:00:00-07:00')] - ), - ] -) -@mock.patch('modules.videos.channel.lib.today', lambda: local_timezone(datetime(2000, 1, 1, 0, 0, 0))) -def test_spread_by_frequency(channels, expected): - result = lib._spread_by_frequency(channels) - for result, (expected_url, expected_frequency, expected_dt) in zip(result, expected): - assert result['url'] == expected_url - assert result['frequency'] == expected_frequency - assert str(result['next_download']) == expected_dt - def test_search_censored_videos(test_session, simple_channel): for i in map(str, range(50)): @@ -237,6 +170,12 @@ def test_validate_video_exception(test_session, simple_channel, video_factory): ('foo .mp4', (None, None, None, 'foo')), ('NA_20000303_vp91w5_Bob's Pancakes.mp4', (None, '20000303', 'vp91w5', 'Bob's Pancakes')), ('NA_NA_vp91w5_Bob's Pancakes.mp4', (None, None, 'vp91w5', 'Bob's Pancakes')), + ('Learning Self-Reliance_20170529_p_MzsCFkUPU_Beekeeping 2017 Part 6 - Merging Hives.mp4', + ('Learning Self-Reliance', '20170529', 'p_MzsCFkUPU', 'Beekeeping 2017 Part 6 - Merging Hives')), + ('Learning Self-Reliance_20170529_p_Mzs_Beekeeping 2017 Part 6 - Merging Hives.mp4', + ('Learning Self-Reliance', '20170529', 'p_Mzs', 'Beekeeping 2017 Part 6 - Merging Hives')), + ('Learning Self-Reliance_20170529_p_Beekeeping 2017 Part 6 - Merging Hives.mp4', + ('Learning Self-Reliance', '20170529', None, 'p_Beekeeping 2017 Part 6 - Merging Hives')), ]) def test_parse_video_file_name(file_name, expected): """ diff --git a/wrolpi/common.py b/wrolpi/common.py index c9875ab5..20b82646 100644 --- a/wrolpi/common.py +++ b/wrolpi/common.py @@ -11,11 +11,12 @@ import re import string import tempfile from datetime import datetime, date +from decimal import Decimal from functools import wraps from itertools import islice, filterfalse, tee from multiprocessing import Event, Queue from pathlib import Path -from typing import Union, Callable, Tuple, Dict, Mapping, List, Iterable, Optional +from typing import Union, Callable, Tuple, Dict, Mapping, List, Iterable, Optional, Generator from urllib.parse import urlunsplit, urlparse import yaml @@ -600,3 +601,40 @@ def chdir(directory: Union[pathlib.Path, str, None] = None, with_home: bool = Fa else: del os.environ['HOME'] return + + +ZIG_TYPE = Union[int, float, complex, Decimal, datetime] + + +def zig_zag(low: ZIG_TYPE, high: ZIG_TYPE) -> Generator[ZIG_TYPE, None, None]: + """ + Generate numbers between `low` and `high` that are + spread out evenly. Produces infinite results. + + >>> list(zig_zag(0, 10)) + [0, 5, 2, 7, 1, 3, 6, 8, 0, 1, 3, 4, 5, 6, 8, 9] + >>> list(zig_zag(0, 5)) + [0, 2, 1, 3, 0, 1, 3, 4] + >>> list(zig_zag(50.0, 100.0)) + [50.0, 75.0, 62.5, 87.5, 56.25, 68.75, 81.25, 93.75, 53.125, 59.375, 65.625] + """ + output_type = type(low) + if not isinstance(high, type(low)): + raise ValueError(f'high and low must be same type') + if isinstance(low, datetime) and isinstance(high, datetime): + low, high = low.timestamp(), high.timestamp() + output_type = datetime.fromtimestamp + + # Special thanks to my wife for helping me solve this! :* + results = set() + num = low + divisor = 2 + diff = high - low + while True: + if num not in results: + yield output_type(num) + results.add(num) + num += diff / divisor + if num >= high: + divisor *= 2 + num = low + (diff / divisor) diff --git a/wrolpi/conftest.py b/wrolpi/conftest.py index fede0150..c8d72b87 100644 --- a/wrolpi/conftest.py +++ b/wrolpi/conftest.py @@ -9,7 +9,6 @@ from unittest import mock from uuid import uuid1 import pytest -import pytest_asyncio import yaml from sanic_testing.testing import SanicTestClient from sqlalchemy.engine import Engine, create_engine @@ -52,7 +51,7 @@ def test_db() -> Tuple[Engine, Session]: return test_engine, session -@pytest_asyncio.fixture() +@pytest.fixture() def test_session() -> Session: """ Pytest Fixture to get a test database session. @@ -75,7 +74,7 @@ def test_session() -> Session: conn.close() -@pytest_asyncio.fixture(autouse=True) +@pytest.fixture(autouse=True) def test_directory() -> pathlib.Path: """ Overwrite the media directory with a temporary directory. @@ -87,7 +86,7 @@ def test_directory() -> pathlib.Path: yield tmp_path -@pytest_asyncio.fixture(autouse=True) +@pytest.fixture(autouse=True) def test_config(): """ Create a test config based off the example config. diff --git a/wrolpi/dates.py b/wrolpi/dates.py index 09dc36f6..14926e7e 100644 --- a/wrolpi/dates.py +++ b/wrolpi/dates.py @@ -1,5 +1,6 @@ import logging from datetime import datetime, date, timezone +from enum import Enum import pytest import pytz @@ -14,6 +15,13 @@ DEFAULT_TIMEZONE = pytz.timezone(DEFAULT_TIMEZONE_STR) TEST_DATETIME: datetime = None +class Seconds(int, Enum): + minute = 60 + hour = minute * 60 + day = hour * 24 + week = day * 7 + + @pytest.fixture def fake_now(): def _(dt: datetime): diff --git a/wrolpi/downloader.py b/wrolpi/downloader.py index 0d5ab963..648a610c 100644 --- a/wrolpi/downloader.py +++ b/wrolpi/downloader.py @@ -3,7 +3,7 @@ import multiprocessing import pathlib import subprocess from collections import defaultdict -from datetime import timedelta +from datetime import timedelta, datetime from enum import Enum from functools import partial from operator import attrgetter @@ -13,8 +13,8 @@ from sqlalchemy import Column, Integer, String, Text from sqlalchemy.dialects.postgresql import JSONB from sqlalchemy.orm import Session -from wrolpi.common import Base, ModelHelper, logger, iterify, wrol_mode_check -from wrolpi.dates import TZDateTime, now +from wrolpi.common import Base, ModelHelper, logger, iterify, wrol_mode_check, zig_zag +from wrolpi.dates import TZDateTime, now, Seconds, local_timezone from wrolpi.db import get_db_session, get_db_curs, get_db_context, optional_session from wrolpi.errors import InvalidDownload, UnrecoverableDownloadError from wrolpi.vars import PYTEST @@ -73,13 +73,6 @@ class Download(ModelHelper, Base): ) return d - def increment_next_download(self): - if self.frequency: - delta = timedelta(seconds=self.frequency) - else: - delta = DEFAULT_RETRY_FREQUENCY - self.next_download = now() + delta - def renew(self, reset_attempts: bool = False): """ Mark this Download as "new" so it will be retried. @@ -93,7 +86,6 @@ class Download(ModelHelper, Base): Download should be tried again after a time. """ self.status = 'deferred' - self.increment_next_download() def fail(self): """ @@ -116,8 +108,6 @@ class Download(ModelHelper, Base): """ self.status = 'complete' self.last_successful_download = now() - if self.frequency: - self.increment_next_download() def get_downloader(self): if self.downloader: @@ -384,8 +374,10 @@ class DownloadManager: download.fail() elif success: download.complete() + download.next_download = self.get_next_download(download, session) else: download.defer() + download.next_download = self.get_next_download(download, session) session_.commit() if download_count: @@ -594,6 +586,9 @@ class DownloadManager: @optional_session def get_fe_downloads(self, session: Session = None): + """ + Get downloads for the Frontend. + """ data = dict( recurring_downloads=self.get_recurring_downloads(session=session), once_downloads=self.get_once_downloads(session=session, limit=100), @@ -604,6 +599,46 @@ class DownloadManager: def get_pending_downloads(self, session: Session): return session.query(Download).filter_by(status='pending').all() + @optional_session + def get_next_download(self, download: Download, session: Session) -> Optional[datetime]: + """ + If the download is "deferred", download soon. But, slowly increase the time between attempts. + + If the download is "complete" and the download has a frequency, schedule the download in it's next iteration. + (Next week, month, etc.) + """ + if download.status == 'deferred': + # Increase next_download slowly at first, then by large gaps later. The largest gap is the download + # frequency. + hours = 3 ** (download.attempts or 1) + seconds = hours * Seconds.hour + if download.frequency: + seconds = min(seconds, download.frequency) + delta = timedelta(seconds=seconds) + return now() + delta + + if not download.frequency: + return None + if not download.id: + raise ValueError('Cannot get next download when Download is not in DB.') + + freq = download.frequency + + # Keep this Download in it's position within the like-frequency downloads. + downloads = [i.id for i in session.query(Download).filter_by(frequency=freq).order_by(Download.id)] + index = downloads.index(download.id) + + # Download was successful. Spread the same-frequency downloads out. + start_date = local_timezone(datetime(2000, 1, 1)) + # Weeks/months/etc since start_date. + iterations = ((now() - start_date) // freq).total_seconds() + start_date = start_date + timedelta(seconds=iterations * freq) + end_date = start_date + timedelta(seconds=freq) + zagger = zig_zag(start_date, end_date) + next_download = [i for i, j in zip(zagger, range(len(downloads)))][index] + next_download = local_timezone(next_download) + return next_download + # The global DownloadManager. This should be used everywhere! download_manager = DownloadManager() diff --git a/wrolpi/test/test_common.py b/wrolpi/test/test_common.py index c9c4a103..c09164b5 100644 --- a/wrolpi/test/test_common.py +++ b/wrolpi/test/test_common.py @@ -3,11 +3,12 @@ import pathlib import tempfile import unittest from datetime import date, datetime +from decimal import Decimal import pytest from sanic_openapi import doc -from wrolpi.common import combine_dicts, insert_parameter, date_range, api_param_limiter, chdir +from wrolpi.common import combine_dicts, insert_parameter, date_range, api_param_limiter, chdir, zig_zag from wrolpi.dates import set_timezone, now from wrolpi.errors import NoBodyContents, MissingRequiredField, ExcessJSONFields, InvalidTimezone from wrolpi.schema import validate_data @@ -346,3 +347,34 @@ def test_chdir(): with chdir(pathlib.Path(d)): assert os.getcwd() == d assert os.environ['HOME'] != os.getcwd() + + +@pytest.mark.parametrize('low,high,expected', [ + (0, 10000, [0, 5000, 2500, 7500, 1250, 3750, 6250, 8750, 625, 1875, 3125, 4375, 5625, 6875, 8125, 9375]), + (0, 1000, [0, 500, 250, 750, 125, 375, 625, 875, 62, 187, 312, 437, 562, 687, 812, 937]), + (0, 100, [0, 50, 25, 75, 12, 37, 62, 87, 6, 18, 31, 43, 56, 68, 81, 93]), + # Values repeat when there are not enough. + (0, 10, [0, 5, 2, 7, 1, 3, 6, 8, 0, 1, 3, 4, 5, 6, 8, 9]), + (0, 5, [0, 2, 1, 3, 0, 1, 3, 4]), + (0, 2, [0, 1, 0, 1, 0, 0, 1, 1, 0, 0, 0, 0, 1, 1, 1, 1]), + (50, 100, [50, 75, 62, 87, 56, 68, 81, 93, 53, 59, 65]), + (8, 98, [8, 53, 30, 75, 19, 41, 64]), + # Floats are supported. Output is identical to above. + (50.0, 100.0, [50.0, 75.0, 62.5, 87.5, 56.25, 68.75, 81.25, 93.75, 53.125, 59.375, 65.625]), + # Decimals are supported. Output is identical to above. + (Decimal('50.0'), Decimal('100.0'), [ + Decimal('50'), Decimal('75'), Decimal('62.5'), Decimal('87.5'), Decimal('56.25'), Decimal('68.75'), + Decimal('81.25'), Decimal('93.75'), Decimal('53.125'), Decimal('59.375'), Decimal('65.625') + ]), + # Datetimes are supported. + (datetime(2000, 1, 1), datetime(2000, 1, 8), [ + datetime(2000, 1, 1), datetime(2000, 1, 4, 12), datetime(2000, 1, 2, 18), datetime(2000, 1, 6, 6), + datetime(2000, 1, 1, 21), datetime(2000, 1, 3, 15), + ]) +]) +def test_zig_zag(low, high, expected): + zagger = zig_zag(low, high) + for i in expected: + result = next(zagger) + assert result == i + assert low <= result < high diff --git a/wrolpi/test/test_downloader.py b/wrolpi/test/test_downloader.py index f575bba5..57d4f7c7 100644 --- a/wrolpi/test/test_downloader.py +++ b/wrolpi/test/test_downloader.py @@ -1,5 +1,5 @@ from abc import ABC -from datetime import datetime, timedelta +from datetime import datetime from itertools import zip_longest from unittest import mock from unittest.mock import MagicMock @@ -9,7 +9,7 @@ import pytest from wrolpi import downloader from wrolpi.dates import local_timezone from wrolpi.db import get_db_context -from wrolpi.downloader import Downloader, Download +from wrolpi.downloader import Downloader, Download, DownloadFrequency from wrolpi.errors import UnrecoverableDownloadError, InvalidDownload, WROLModeEnabled from wrolpi.test.common import wrap_test_db, TestAPI @@ -185,7 +185,7 @@ class TestDownloader(TestAPI): downloads = list(self.mgr.get_recurring_downloads(session)) self.assertEqual(len(downloads), 1) download = downloads[0] - expected = local_timezone(datetime(2020, 1, 1, 1, 0, 0)) + expected = local_timezone(datetime(2020, 1, 1, 0, 0, 0)) self.assertEqual(download.next_download, expected) self.assertEqual(download.last_successful_download, now) @@ -195,7 +195,7 @@ class TestDownloader(TestAPI): self.assertEqual(download.last_successful_download, now) # Download is due an hour later. - mock_now.return_value = local_timezone(datetime(2020, 1, 1, 1, 0, 1)) + mock_now.return_value = local_timezone(datetime(2020, 1, 1, 2, 0, 1)) self.mgr.renew_recurring_downloads(session) downloads = list(self.mgr.get_new_downloads(session)) self.assertEqual(len(downloads), 1) @@ -214,11 +214,12 @@ class TestDownloader(TestAPI): self.assertEqual(download.status, 'deferred') self.assertEqual(download.last_successful_download, now) # Download should be retried after the DEFAULT_RETRY_FREQUENCY. - self.assertEqual(download.next_download, expected + timedelta(hours=1, seconds=1)) + expected = local_timezone(datetime(2020, 1, 1, 3, 0, 1)) + self.assertEqual(download.next_download, expected) # Try the download again, it finally succeeds. http_downloader.do_download.reset_mock() - now = local_timezone(datetime(2020, 1, 1, 3, 0, 0)) + now = local_timezone(datetime(2020, 1, 1, 4, 0, 1)) mock_now.return_value = now http_downloader.do_download.return_value = True self.mgr.renew_recurring_downloads(session) @@ -227,7 +228,8 @@ class TestDownloader(TestAPI): download = session.query(Download).one() self.assertEqual(download.status, 'complete') self.assertEqual(download.last_successful_download, now) - self.assertEqual(download.next_download, local_timezone(datetime(2020, 1, 1, 4, 0, 0))) + # Floats cause slightly wrong date. + self.assertEqual(download.next_download, local_timezone(datetime(2020, 1, 1, 4, 0, 0, 997200))) @wrap_test_db @mock.patch('wrolpi.downloader.now', lambda: local_timezone(datetime(2020, 6, 5, 0, 0))) @@ -333,3 +335,34 @@ def test_download_get_downloader(test_session, test_download_manager): download2 = test_download_manager.create_download('https://example.com', test_session, skip_download=True) download2.downloader = 'bad downloader' assert download2.get_downloader() is None + + +def test_get_next_download(test_session, test_download_manager, fake_now): + fake_now(datetime(2000, 1, 1)) + download = Download() + download.frequency = DownloadFrequency.weekly + download.status = 'deferred' + + # next_download slowly increases as we accumulate attempts. Largest gap is the download frequency. + attempts_expected = [ + (0, local_timezone(datetime(2000, 1, 1, 3))), + (1, local_timezone(datetime(2000, 1, 1, 3))), + (2, local_timezone(datetime(2000, 1, 1, 9))), + (3, local_timezone(datetime(2000, 1, 2, 3))), + (4, local_timezone(datetime(2000, 1, 4, 9))), + (5, local_timezone(datetime(2000, 1, 8))), + (6, local_timezone(datetime(2000, 1, 8))), + ] + for attempts, expected in attempts_expected: + download.attempts = attempts + result = test_download_manager.get_next_download(download) + assert result == expected, f'{attempts} != {result}' + + d1 = Download(url='https://example.com/1', frequency=DownloadFrequency.weekly) + d2 = Download(url='https://example.com/2', frequency=DownloadFrequency.weekly) + d3 = Download(url='https://example.com/3', frequency=DownloadFrequency.weekly) + test_session.add_all([d1, d2, d3]) + test_session.commit() + assert test_download_manager.get_next_download(d1) == local_timezone(datetime(2000, 1, 1)) + assert test_download_manager.get_next_download(d2) == local_timezone(datetime(2000, 1, 4, 12)) + assert test_download_manager.get_next_download(d3) == local_timezone(datetime(2000, 1, 2, 18))