Skip to content

Commit

Permalink
unit tests for loading dropins
Browse files Browse the repository at this point in the history
  • Loading branch information
pjrobertson committed Jan 21, 2025
1 parent 9dde9b2 commit 6814205
Show file tree
Hide file tree
Showing 3 changed files with 28 additions and 10 deletions.
14 changes: 7 additions & 7 deletions src/auto_archiver/archivers/generic_archiver/generic_archiver.py
Original file line number Diff line number Diff line change
Expand Up @@ -79,7 +79,7 @@ def download_additional_media(self, video_data: dict, info_extractor: InfoExtrac
except Exception as e:
logger.error(f"Error downloading cover image {thumbnail_url}: {e}")

dropin = self.dropin_for_extractor(info_extractor)
dropin = self.dropin_for_name(info_extractor.ie_key())
if dropin:
try:
metadata = dropin.download_additional_media(video_data, info_extractor, metadata)
Expand All @@ -102,7 +102,7 @@ def keys_to_clean(self, info_extractor: InfoExtractor, video_data: dict) -> dict
'_format_sort_fields', 'chapters', 'requested_formats', 'format_note',
'audio_channels', 'asr', 'fps', 'was_live', 'is_live', 'heatmap', 'age_limit', 'stretched_ratio']

dropin = self.dropin_for_extractor(info_extractor)
dropin = self.dropin_for_name(info_extractor.ie_key())
if dropin:
try:
base_keys += dropin.keys_to_clean(video_data, info_extractor)
Expand Down Expand Up @@ -157,7 +157,7 @@ def get_metatdata_for_post(self, info_extractor: Type[InfoExtractor], url: str,

ie_instance = info_extractor(downloader=ydl)
post_data = None
dropin = self.dropin_for_extractor(info_extractor)
dropin = self.dropin_for_name(info_extractor.ie_key())
if not dropin:
# TODO: add a proper link to 'how to create your own dropin'
logger.debug(f"""Could not find valid dropin for {info_extractor.IE_NAME}.
Expand Down Expand Up @@ -207,8 +207,7 @@ def get_metadata_for_video(self, data: dict, info_extractor: Type[InfoExtractor]

return self.add_metadata(data, info_extractor, url, result)

def dropin_for_extractor(self, info_extractor: Type[InfoExtractor], additional_paths = []):
dropin_name = info_extractor.ie_key().lower()
def dropin_for_name(self, dropin_name: str, additional_paths = [], package=__package__) -> Type[InfoExtractor]:

if dropin_name == "generic":
# no need for a dropin for the generic extractor (?)
Expand All @@ -228,6 +227,7 @@ def _load_dropin(dropin):
# which would allow the user to override the default dropins/add their own
paths = [] + additional_paths
for path in paths:
breakpoint()
dropin_path = os.path.join(path, f"{dropin_name}.py")
dropin_spec = importlib.util.spec_from_file_location(dropin_name, dropin_path)
if not dropin_spec:
Expand All @@ -241,7 +241,7 @@ def _load_dropin(dropin):

# fallback to loading the dropins within auto-archiver
try:
return _load_dropin(importlib.import_module(f".{dropin_name}", package=__package__))
return _load_dropin(importlib.import_module(f".{dropin_name}", package=package))
except ModuleNotFoundError:
pass

Expand All @@ -258,7 +258,7 @@ def download_for_extractor(self, info_extractor: InfoExtractor, url: str, ydl: y
ydl.params['getcomments'] = False
result = False

dropin_submodule = self.dropin_for_extractor(info_extractor)
dropin_submodule = self.dropin_for_name(info_extractor.ie_key())

try:
if dropin_submodule and dropin_submodule.skip_ytdlp_download(info_extractor, url):
Expand Down
19 changes: 16 additions & 3 deletions tests/archivers/test_generic_archiver.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,12 @@
import pytest
from pathlib import Path
import datetime
import datetime
import os

from auto_archiver.archivers.generic_archiver import GenericArchiver
from os.path import dirname

import pytest

from auto_archiver.archivers.generic_archiver import GenericArchiver
from .test_archiver_base import TestArchiverBase

class TestGenericArchiver(TestArchiverBase):
Expand All @@ -23,6 +25,17 @@ class TestGenericArchiver(TestArchiverBase):
'cookies_from_browser': False,
'cookie_file': None,
}

def test_load_dropin(self):
# test loading dropins that are in the generic_archiver package
package = "auto_archiver.archivers.generic_archiver"
assert self.archiver.dropin_for_name("bluesky", package=package)

# test loading dropings via filepath
path = os.path.join(dirname(dirname(__file__)), "data/")
assert self.archiver.dropin_for_name("dropin", additional_paths=[path])



@pytest.mark.parametrize("url, is_suitable", [
("https://www.youtube.com/watch?v=5qap5aO4i9A", True),
Expand Down
5 changes: 5 additions & 0 deletions tests/data/dropin.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
# this is a dummy class used to test importing a dropin in the
# generic extractor by filename/path

class Dropin:
pass

0 comments on commit 6814205

Please sign in to comment.