From 801784c46d7764832f98e988faa804e199003a23 Mon Sep 17 00:00:00 2001 From: Serene-Arc Date: Fri, 5 Nov 2021 13:23:55 +1000 Subject: [PATCH] Fix a crash when downloading a disabled pornhub video --- bdfr/site_downloaders/pornhub.py | 8 +++++++- tests/site_downloaders/test_pornhub.py | 14 ++++++++++++++ 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/bdfr/site_downloaders/pornhub.py b/bdfr/site_downloaders/pornhub.py index c2bc0ad..748454e 100644 --- a/bdfr/site_downloaders/pornhub.py +++ b/bdfr/site_downloaders/pornhub.py @@ -6,6 +6,7 @@ from typing import Optional from praw.models import Submission +from bdfr.exceptions import SiteDownloaderError from bdfr.resource import Resource from bdfr.site_authenticator import SiteAuthenticator from bdfr.site_downloaders.youtube import Youtube @@ -22,10 +23,15 @@ class PornHub(Youtube): 'format': 'best', 'nooverwrites': True, } + if video_attributes := super().get_video_attributes(self.post.url): + extension = video_attributes['ext'] + else: + raise SiteDownloaderError() + out = Resource( self.post, self.post.url, super()._download_video(ytdl_options), - super().get_video_attributes(self.post.url)['ext'], + extension, ) return [out] diff --git a/tests/site_downloaders/test_pornhub.py b/tests/site_downloaders/test_pornhub.py index 5c220cc..cbe3662 100644 --- a/tests/site_downloaders/test_pornhub.py +++ b/tests/site_downloaders/test_pornhub.py @@ -5,6 +5,7 @@ from unittest.mock import MagicMock import pytest +from bdfr.exceptions import SiteDownloaderError from bdfr.resource import Resource from bdfr.site_downloaders.pornhub import PornHub @@ -13,6 +14,7 @@ from bdfr.site_downloaders.pornhub import PornHub @pytest.mark.slow @pytest.mark.parametrize(('test_url', 'expected_hash'), ( ('https://www.pornhub.com/view_video.php?viewkey=ph6074c59798497', 'd9b99e4ebecf2d8d67efe5e70d2acf8a'), + ('https://www.pornhub.com/view_video.php?viewkey=ph5ede121f0d3f8', ''), )) def test_find_resources_good(test_url: str, expected_hash: str): test_submission = MagicMock() @@ -23,3 +25,15 @@ def test_find_resources_good(test_url: str, expected_hash: str): assert isinstance(resources[0], Resource) resources[0].download() assert resources[0].hash.hexdigest() == expected_hash + + +@pytest.mark.online +@pytest.mark.parametrize('test_url', ( + 'https://www.pornhub.com/view_video.php?viewkey=ph5ede121f0d3f8', +)) +def test_find_resources_good(test_url: str): + test_submission = MagicMock() + test_submission.url = test_url + downloader = PornHub(test_submission) + with pytest.raises(SiteDownloaderError): + downloader.find_resources()