diff --git a/server/szurubooru/errors.py b/server/szurubooru/errors.py index acfe8f1c..beeb469c 100644 --- a/server/szurubooru/errors.py +++ b/server/szurubooru/errors.py @@ -56,6 +56,3 @@ class InvalidParameterError(ValidationError): class ThirdPartyError(BaseError): pass - -class DownloadTooLargeError(ProcessingError): - pass diff --git a/server/szurubooru/func/net.py b/server/szurubooru/func/net.py index 882e519f..17f46548 100644 --- a/server/szurubooru/func/net.py +++ b/server/szurubooru/func/net.py @@ -29,36 +29,26 @@ def download(url: str, use_video_downloader: bool = False) -> bytes: def _youtube_dl_wrapper(url: str) -> bytes: + outpath = os.path.join( + config.config['data_dir'], + 'temporary-uploads', + 'youtubedl-' + util.get_sha1(url)[0:8] + '.dat') options = { - 'quiet': True, 'ignoreerrors': False, 'format': 'best[ext=webm]/best[ext=mp4]/best[ext=flv]', 'logger': logger, - 'noplaylist': True, 'max_filesize': config.config['max_dl_filesize'], 'max_downloads': 1, - 'outtmpl': os.path.join( - config.config['data_dir'], - 'temporary-uploads', - 'youtubedl-' + util.get_sha1(url)[0:8] + '.%(ext)s'), + 'outtmpl': outpath, } - with YoutubeDL(options) as ydl: - try: - ydl_info = ydl.extract_info(url, download=True) - # need to confirm if download was skipped due to size - if 'filesize' in ydl_info: - if ydl_info['filesize'] > config.config['max_dl_filesize']: - raise errors.DownloadTooLargeError( - 'Requested video too large (%d MB > %d MB)' % ( - ydl_info['filesize'] / 1.0e6, - config.config['max_dl_filesize'] / 1.0e6)) - ydl_filename = ydl.prepare_filename(ydl_info) - except YoutubeDLError as ex: - raise errors.ThirdPartyError( - 'Error downloading video %s (%s)' % (url, ex)) try: - with open(ydl_filename, 'rb') as f: + with YoutubeDL(options) as ydl: + ydl.extract_info(url, download=True) + with open(outpath, 'rb') as f: return f.read() + except YoutubeDLError as ex: + raise errors.ThirdPartyError( + 'Error downloading video %s (%s)' % (url, ex)) except FileNotFoundError as ex: raise errors.ThirdPartyError( - 'Error downloading video %s' % (url)) + 'Error downloading video %s (file could not be saved)' % (url)) diff --git a/server/szurubooru/tests/func/test_net.py b/server/szurubooru/tests/func/test_net.py index 64951e23..7fe14b62 100644 --- a/server/szurubooru/tests/func/test_net.py +++ b/server/szurubooru/tests/func/test_net.py @@ -58,31 +58,30 @@ def test_download(): assert actual_content == expected_content -def test_too_large_download(): +@pytest.mark.parametrize('url', [ + 'https://samples.ffmpeg.org/MPEG-4/video.mp4', +]) +def test_too_large_download(url): pytest.xfail('Download limit not implemented yet') - url = 'https://samples.ffmpeg.org/MPEG-4/video.mp4' - - with pytest.raises(errors.DownloadTooLargeError): + with pytest.raises(errors.ProcessingError): net.download(url) -def test_video_download(): - url = 'https://www.youtube.com/watch?v=C0DPdy98e4c' - expected_sha1 = '365af1c8f59c6865e1a84c6e13e3e25ff89e0ba1' - +@pytest.mark.parametrize('url,expected_sha1', [ + ('https://www.youtube.com/watch?v=C0DPdy98e4c', + '365af1c8f59c6865e1a84c6e13e3e25ff89e0ba1'), + ('https://gfycat.com/immaterialchillyiberianmole', + '953000e81d7bd1da95ce264f872e7b6c4a6484be'), +]) +def test_video_download(url, expected_sha1): actual_content = net.download(url, use_video_downloader=True) assert get_sha1(actual_content) == expected_sha1 -def test_failed_video_download(): - url = 'https://samples.ffmpeg.org/flac/short.flac' - +@pytest.mark.parametrize('url', [ + 'https://samples.ffmpeg.org/flac/short.flac', # not a video + 'https://www.youtube.com/watch?v=dQw4w9WgXcQ', # video too large +]) +def test_failed_video_download(url): with pytest.raises(errors.ThirdPartyError): net.download(url, use_video_downloader=True) - - -def test_too_large_video_download(): - url = 'https://www.youtube.com/watch?v=dQw4w9WgXcQ' - - with pytest.raises(errors.DownloadTooLargeError): - net.download(url, use_video_downloader=True)