server/posts/upload: refactor youtube-dl caller code to fix some bugs

3 files changed, 29 insertions, 43 deletions

diff --git a/server/szurubooru/errors.py b/server/szurubooru/errors.py
index acfe8f1..beeb469 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 882e519..17f4654 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 64951e2..7fe14b6 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)