diff --git a/bdfr/site_downloaders/erome.py b/bdfr/site_downloaders/erome.py index 6130560..6250415 100644 --- a/bdfr/site_downloaders/erome.py +++ b/bdfr/site_downloaders/erome.py @@ -2,7 +2,7 @@ import logging import re -from typing import Optional +from typing import Callable, Optional import bs4 from praw.models import Submission @@ -29,7 +29,7 @@ class Erome(BaseDownloader): for link in links: if not re.match(r'https?://.*', link): link = 'https://' + link - out.append(Resource(self.post, link, Resource.retry_download(link))) + out.append(Resource(self.post, link, self.erome_download(link))) return out @staticmethod @@ -43,3 +43,14 @@ class Erome(BaseDownloader): out.extend([vid.get('src') for vid in videos]) return set(out) + + @staticmethod + def erome_download(url: str) -> Callable: + download_parameters = { + 'headers': { + 'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko)' + ' Chrome/88.0.4324.104 Safari/537.36', + 'Referer': 'https://www.erome.com/', + }, + } + return lambda global_params: Resource.http_download(url, global_params | download_parameters) diff --git a/tests/site_downloaders/test_erome.py b/tests/site_downloaders/test_erome.py index bab34bb..d27cbb7 100644 --- a/tests/site_downloaders/test_erome.py +++ b/tests/site_downloaders/test_erome.py @@ -30,25 +30,18 @@ def test_get_link(test_url: str, expected_urls: tuple[str]): @pytest.mark.online @pytest.mark.slow -@pytest.mark.parametrize(('test_url', 'expected_hashes'), ( - ('https://www.erome.com/a/vqtPuLXh', { - '5da2a8d60d87bed279431fdec8e7d72f' - }), - ('https://www.erome.com/a/lGrcFxmb', { - '0e98f9f527a911dcedde4f846bb5b69f', - '25696ae364750a5303fc7d7dc78b35c1', - '63775689f438bd393cde7db6d46187de', - 'a1abf398cfd4ef9cfaf093ceb10c746a', - 'bd9e1a4ea5ef0d6ba47fb90e337c2d14' - }), +@pytest.mark.parametrize(('test_url', 'expected_hashes_len'), ( + ('https://www.erome.com/a/vqtPuLXh', 1), + ('https://www.erome.com/a/4tP3KI6F', 1), )) -def test_download_resource(test_url: str, expected_hashes: tuple[str]): +def test_download_resource(test_url: str, expected_hashes_len: int): # Can't compare hashes for this test, Erome doesn't return the exact same file from request to request so the hash # will change back and forth randomly mock_submission = MagicMock() mock_submission.url = test_url test_site = Erome(mock_submission) resources = test_site.find_resources() - [res.download() for res in resources] + for res in resources: + res.download() resource_hashes = [res.hash.hexdigest() for res in resources] - assert len(resource_hashes) == len(expected_hashes) + assert len(resource_hashes) == expected_hashes_len