Pele, regarded by many as football\'s greatest ever player, has died in Brazil at the age of 82. https://www.bbc.com/sport/football/42751517
', 'account': {'id': '107905163010312793', 'username': 'bbcnewa', 'acct': 'bbcnewa', 'display_name': 'BBC News', 'locked': False, 'bot': False, 'discoverable': True, 'group': False, 'created_at': '2022-03-05T17:42:01.159Z', 'note': 'News, features and analysis by the BBC
', 'url': 'https://truthsocial.com/@bbcnewa', 'avatar': 'https://static-assets-1.truthsocial.com/tmtg:prime-ts-assets/accounts/avatars/107/905/163/010/312/793/original/e7c07550dc22c23a.jpeg', 'avatar_static': 'https://static-assets-1.truthsocial.com/tmtg:prime-ts-assets/accounts/avatars/107/905/163/010/312/793/original/e7c07550dc22c23a.jpeg', 'header': 'https://static-assets-1.truthsocial.com/tmtg:prime-ts-assets/accounts/headers/107/905/163/010/312/793/original/a00eeec2b57206c7.jpeg', 'header_static': 'https://static-assets-1.truthsocial.com/tmtg:prime-ts-assets/accounts/headers/107/905/163/010/312/793/original/a00eeec2b57206c7.jpeg', 'followers_count': 1131, 'following_count': 3, 'statuses_count': 9, 'last_status_at': '2024-11-12', 'verified': False, 'location': '', 'website': 'https://www.bbc.com/news', 'unauth_visibility': True, 'chats_onboarded': True, 'feeds_onboarded': True, 'accepting_messages': False, 'show_nonmember_group_statuses': None, 'emojis': [], 'fields': [], 'tv_onboarded': True, 'tv_account': False}, 'media_attachments': [], 'mentions': [], 'tags': [], 'card': None, 'group': None, 'quote': None, 'in_reply_to': None, 'reblog': None, 'sponsored': False, 'replies_count': 1, 'reblogs_count': 0, 'favourites_count': 2, 'favourited': False, 'reblogged': False, 'muted': False, 'pinned': False, 'bookmarked': False, 'poll': None, 'emojis': []} """ result = Metadata() result.set_url(url) - timestamp = post['created_at'] # format is 2022-12-29T19:51:18.161Z + timestamp = post["created_at"] # format is 2022-12-29T19:51:18.161Z result.set_timestamp(parse_dt(timestamp)) - result.set('description', post['content']) - result.set('author', post['account']['username']) + result.set("description", post["content"]) + result.set("author", post["account"]["username"]) - for key in ['replies_count', 'reblogs_count', 'favourites_count', ('account', 'followers_count'), ('account', 'following_count'), ('account', 'statuses_count'), ('account', 'display_name'), 'language', 'in_reply_to_account', 'replies_count']: + for key in [ + "replies_count", + "reblogs_count", + "favourites_count", + ("account", "followers_count"), + ("account", "following_count"), + ("account", "statuses_count"), + ("account", "display_name"), + "language", + "in_reply_to_account", + "replies_count", + ]: if isinstance(key, tuple): store_key = " ".join(key) else: store_key = key result.set(store_key, traverse_obj(post, key)) - - # add the media - for media in post.get('media_attachments', []): - filename = archiver.download_from_url(media['url']) - result.add_media(Media(filename), id=media.get('id')) - return result \ No newline at end of file + # add the media + for media in post.get("media_attachments", []): + filename = archiver.download_from_url(media["url"]) + result.add_media(Media(filename), id=media.get("id")) + + return result diff --git a/src/auto_archiver/modules/generic_extractor/twitter.py b/src/auto_archiver/modules/generic_extractor/twitter.py index 3faed6b..e4cbe74 100644 --- a/src/auto_archiver/modules/generic_extractor/twitter.py +++ b/src/auto_archiver/modules/generic_extractor/twitter.py @@ -1,4 +1,6 @@ -import re, mimetypes, json +import re +import mimetypes +import json from datetime import datetime from loguru import logger @@ -10,9 +12,8 @@ from auto_archiver.core.extractor import Extractor from .dropin import GenericDropin, InfoExtractor + class Twitter(GenericDropin): - - def choose_variant(self, variants): # choosing the highest quality possible variant, width, height = None, 0, 0 @@ -27,44 +28,43 @@ class Twitter(GenericDropin): else: variant = var if not variant else variant return variant - + def extract_post(self, url: str, ie_instance: InfoExtractor): - twid = ie_instance._match_valid_url(url).group('id') + twid = ie_instance._match_valid_url(url).group("id") return ie_instance._extract_status(twid=twid) def create_metadata(self, tweet: dict, ie_instance: InfoExtractor, archiver: Extractor, url: str) -> Metadata: result = Metadata() try: if not tweet.get("user") or not tweet.get("created_at"): - raise ValueError(f"Error retreiving post. Are you sure it exists?") + raise ValueError("Error retreiving post. Are you sure it exists?") timestamp = datetime.strptime(tweet["created_at"], "%a %b %d %H:%M:%S %z %Y") except (ValueError, KeyError) as ex: logger.warning(f"Unable to parse tweet: {str(ex)}\nRetreived tweet data: {tweet}") return False - - result\ - .set_title(tweet.get('full_text', ''))\ - .set_content(json.dumps(tweet, ensure_ascii=False))\ - .set_timestamp(timestamp) + + result.set_title(tweet.get("full_text", "")).set_content(json.dumps(tweet, ensure_ascii=False)).set_timestamp( + timestamp + ) if not tweet.get("entities", {}).get("media"): - logger.debug('No media found, archiving tweet text only') + logger.debug("No media found, archiving tweet text only") result.status = "twitter-ytdl" return result for i, tw_media in enumerate(tweet["entities"]["media"]): media = Media(filename="") mimetype = "" if tw_media["type"] == "photo": - media.set("src", UrlUtil.twitter_best_quality_url(tw_media['media_url_https'])) + media.set("src", UrlUtil.twitter_best_quality_url(tw_media["media_url_https"])) mimetype = "image/jpeg" elif tw_media["type"] == "video": - variant = self.choose_variant(tw_media['video_info']['variants']) - media.set("src", variant['url']) - mimetype = variant['content_type'] + variant = self.choose_variant(tw_media["video_info"]["variants"]) + media.set("src", variant["url"]) + mimetype = variant["content_type"] elif tw_media["type"] == "animated_gif": - variant = tw_media['video_info']['variants'][0] - media.set("src", variant['url']) - mimetype = variant['content_type'] + variant = tw_media["video_info"]["variants"][0] + media.set("src", variant["url"]) + mimetype = variant["content_type"] ext = mimetypes.guess_extension(mimetype) - media.filename = archiver.download_from_url(media.get("src"), f'{slugify(url)}_{i}{ext}') + media.filename = archiver.download_from_url(media.get("src"), f"{slugify(url)}_{i}{ext}") result.add_media(media) - return result \ No newline at end of file + return result diff --git a/src/auto_archiver/modules/gsheet_feeder_db/__init__.py b/src/auto_archiver/modules/gsheet_feeder_db/__init__.py index 2e9ac02..fbd37b9 100644 --- a/src/auto_archiver/modules/gsheet_feeder_db/__init__.py +++ b/src/auto_archiver/modules/gsheet_feeder_db/__init__.py @@ -1,2 +1,2 @@ from .gworksheet import GWorksheet -from .gsheet_feeder_db import GsheetsFeederDB \ No newline at end of file +from .gsheet_feeder_db import GsheetsFeederDB diff --git a/src/auto_archiver/modules/gsheet_feeder_db/__manifest__.py b/src/auto_archiver/modules/gsheet_feeder_db/__manifest__.py index 9dd6b87..5143218 100644 --- a/src/auto_archiver/modules/gsheet_feeder_db/__manifest__.py +++ b/src/auto_archiver/modules/gsheet_feeder_db/__manifest__.py @@ -12,9 +12,7 @@ "default": None, "help": "the id of the sheet to archive (alternative to 'sheet' config)", }, - "header": {"default": 1, - "type": "int", - "help": "index of the header row (starts at 1)", "type": "int"}, + "header": {"default": 1, "help": "index of the header row (starts at 1)", "type": "int"}, "service_account": { "default": "secrets/service_account.json", "help": "service account JSON file path. Learn how to create one: https://gspread.readthedocs.io/en/latest/oauth2.html", @@ -53,19 +51,6 @@ "help": "if True the stored files path will include 'workbook_name/worksheet_name/...'", "type": "bool", }, - "allow_worksheets": { - "default": set(), - "help": "(CSV) only worksheets whose name is included in allow are included (overrides worksheet_block), leave empty so all are allowed", - }, - "block_worksheets": { - "default": set(), - "help": "(CSV) explicitly block some worksheets from being processed", - }, - "use_sheet_names_in_stored_paths": { - "default": True, - "type": "bool", - "help": "if True the stored files path will include 'workbook_name/worksheet_name/...'", - } }, "description": """ GsheetsFeederDatabase diff --git a/src/auto_archiver/modules/gsheet_feeder_db/gsheet_feeder_db.py b/src/auto_archiver/modules/gsheet_feeder_db/gsheet_feeder_db.py index 406eeb4..109be3f 100644 --- a/src/auto_archiver/modules/gsheet_feeder_db/gsheet_feeder_db.py +++ b/src/auto_archiver/modules/gsheet_feeder_db/gsheet_feeder_db.py @@ -8,6 +8,7 @@ The filtered rows are processed into `Metadata` objects. - validates the sheet's structure and filters rows based on input configurations. - Ensures only rows with valid URLs and unprocessed statuses are included. """ + import os from typing import Tuple, Union from urllib.parse import quote @@ -19,11 +20,10 @@ from slugify import slugify from auto_archiver.core import Feeder, Database, Media from auto_archiver.core import Metadata from auto_archiver.modules.gsheet_feeder_db import GWorksheet -from auto_archiver.utils.misc import calculate_file_hash, get_current_timestamp +from auto_archiver.utils.misc import get_current_timestamp class GsheetsFeederDB(Feeder, Database): - def setup(self) -> None: self.gsheets_client = gspread.service_account(filename=self.service_account) # TODO mv to validators @@ -42,24 +42,28 @@ class GsheetsFeederDB(Feeder, Database): if not self.should_process_sheet(worksheet.title): logger.debug(f"SKIPPED worksheet '{worksheet.title}' due to allow/block rules") continue - logger.info(f'Opening worksheet {ii=}: {worksheet.title=} header={self.header}') + logger.info(f"Opening worksheet {ii=}: {worksheet.title=} header={self.header}") gw = GWorksheet(worksheet, header_row=self.header, columns=self.columns) if len(missing_cols := self.missing_required_columns(gw)): - logger.warning(f"SKIPPED worksheet '{worksheet.title}' due to missing required column(s) for {missing_cols}") + logger.warning( + f"SKIPPED worksheet '{worksheet.title}' due to missing required column(s) for {missing_cols}" + ) continue # process and yield metadata here: yield from self._process_rows(gw) - logger.success(f'Finished worksheet {worksheet.title}') + logger.success(f"Finished worksheet {worksheet.title}") def _process_rows(self, gw: GWorksheet): for row in range(1 + self.header, gw.count_rows() + 1): - url = gw.get_cell(row, 'url').strip() - if not len(url): continue - original_status = gw.get_cell(row, 'status') - status = gw.get_cell(row, 'status', fresh=original_status in ['', None]) + url = gw.get_cell(row, "url").strip() + if not len(url): + continue + original_status = gw.get_cell(row, "status") + status = gw.get_cell(row, "status", fresh=original_status in ["", None]) # TODO: custom status parser(?) aka should_retry_from_status - if status not in ['', None]: continue + if status not in ["", None]: + continue # All checks done - archival process starts here m = Metadata().set_url(url) @@ -70,10 +74,10 @@ class GsheetsFeederDB(Feeder, Database): # TODO: Check folder value not being recognised m.set_context("gsheet", {"row": row, "worksheet": gw}) - if gw.get_cell_or_default(row, 'folder', "") is None: - folder = '' + if gw.get_cell_or_default(row, "folder", "") is None: + folder = "" else: - folder = slugify(gw.get_cell_or_default(row, 'folder', "").strip()) + folder = slugify(gw.get_cell_or_default(row, "folder", "").strip()) if len(folder): if self.use_sheet_names_in_stored_paths: m.set_context("folder", os.path.join(folder, slugify(self.sheet), slugify(gw.wks.title))) @@ -91,12 +95,11 @@ class GsheetsFeederDB(Feeder, Database): def missing_required_columns(self, gw: GWorksheet) -> list: missing = [] - for required_col in ['url', 'status']: + for required_col in ["url", "status"]: if not gw.col_exists(required_col): missing.append(required_col) return missing - def started(self, item: Metadata) -> None: logger.warning(f"STARTED {item}") gw, row = self._retrieve_gsheet(item) @@ -155,9 +158,7 @@ class GsheetsFeederDB(Feeder, Database): if len(pdq_hashes): batch_if_valid("pdq_hash", ",".join(pdq_hashes)) - if (screenshot := item.get_media_by_id("screenshot")) and hasattr( - screenshot, "urls" - ): + if (screenshot := item.get_media_by_id("screenshot")) and hasattr(screenshot, "urls"): batch_if_valid("screenshot", "\n".join(screenshot.urls)) if thumbnail := item.get_first_image("thumbnail"): @@ -186,11 +187,12 @@ class GsheetsFeederDB(Feeder, Database): logger.debug(f"Unable to update sheet: {e}") def _retrieve_gsheet(self, item: Metadata) -> Tuple[GWorksheet, int]: - if gsheet := item.get_context("gsheet"): gw: GWorksheet = gsheet.get("worksheet") row: int = gsheet.get("row") elif self.sheet_id: - logger.error(f"Unable to retrieve Gsheet for {item.get_url()}, GsheetDB must be used alongside GsheetFeeder.") + logger.error( + f"Unable to retrieve Gsheet for {item.get_url()}, GsheetDB must be used alongside GsheetFeeder." + ) return gw, row diff --git a/src/auto_archiver/modules/gsheet_feeder_db/gworksheet.py b/src/auto_archiver/modules/gsheet_feeder_db/gworksheet.py index ba2d691..6dac059 100644 --- a/src/auto_archiver/modules/gsheet_feeder_db/gworksheet.py +++ b/src/auto_archiver/modules/gsheet_feeder_db/gworksheet.py @@ -5,24 +5,25 @@ class GWorksheet: """ This class makes read/write operations to the a worksheet easier. It can read the headers from a custom row number, but the row references - should always include the offset of the header. - eg: if header=4, row 5 will be the first with data. + should always include the offset of the header. + eg: if header=4, row 5 will be the first with data. """ + COLUMN_NAMES = { - 'url': 'link', - 'status': 'archive status', - 'folder': 'destination folder', - 'archive': 'archive location', - 'date': 'archive date', - 'thumbnail': 'thumbnail', - 'timestamp': 'upload timestamp', - 'title': 'upload title', - 'text': 'text content', - 'screenshot': 'screenshot', - 'hash': 'hash', - 'pdq_hash': 'perceptual hashes', - 'wacz': 'wacz', - 'replaywebpage': 'replaywebpage', + "url": "link", + "status": "archive status", + "folder": "destination folder", + "archive": "archive location", + "date": "archive date", + "thumbnail": "thumbnail", + "timestamp": "upload timestamp", + "title": "upload title", + "text": "text content", + "screenshot": "screenshot", + "hash": "hash", + "pdq_hash": "perceptual hashes", + "wacz": "wacz", + "replaywebpage": "replaywebpage", } def __init__(self, worksheet, columns=COLUMN_NAMES, header_row=1): @@ -36,7 +37,7 @@ class GWorksheet: def _check_col_exists(self, col: str): if col not in self.columns: - raise Exception(f'Column {col} is not in the configured column names: {self.columns.keys()}') + raise Exception(f"Column {col} is not in the configured column names: {self.columns.keys()}") def _col_index(self, col: str): self._check_col_exists(col) @@ -58,7 +59,7 @@ class GWorksheet: def get_cell(self, row, col: str, fresh=False): """ - returns the cell value from (row, col), + returns the cell value from (row, col), where row can be an index (1-based) OR list of values as received from self.get_row(row) if fresh=True, the sheet is queried again for this cell @@ -67,11 +68,11 @@ class GWorksheet: if fresh: return self.wks.cell(row, col_index + 1).value - if type(row) == int: + if isinstance(row, int): row = self.get_row(row) if col_index >= len(row): - return '' + return "" return row[col_index] def get_cell_or_default(self, row, col: str, default: str = None, fresh=False, when_empty_use_default=True): @@ -83,7 +84,7 @@ class GWorksheet: if when_empty_use_default and val.strip() == "": return default return val - except: + except Exception: return default def set_cell(self, row: int, col: str, val): @@ -96,13 +97,9 @@ class GWorksheet: receives a list of [(row:int, col:str, val)] and batch updates it, the parameters are the same as in the self.set_cell() method """ cell_updates = [ - { - 'range': self.to_a1(row, col), - 'values': [[str(val)[0:49999]]] - } - for row, col, val in cell_updates + {"range": self.to_a1(row, col), "values": [[str(val)[0:49999]]]} for row, col, val in cell_updates ] - self.wks.batch_update(cell_updates, value_input_option='USER_ENTERED') + self.wks.batch_update(cell_updates, value_input_option="USER_ENTERED") def to_a1(self, row: int, col: str): # row is 1-based diff --git a/src/auto_archiver/modules/hash_enricher/__init__.py b/src/auto_archiver/modules/hash_enricher/__init__.py index 18ec885..3532e93 100644 --- a/src/auto_archiver/modules/hash_enricher/__init__.py +++ b/src/auto_archiver/modules/hash_enricher/__init__.py @@ -1 +1 @@ -from .hash_enricher import HashEnricher \ No newline at end of file +from .hash_enricher import HashEnricher diff --git a/src/auto_archiver/modules/hash_enricher/__manifest__.py b/src/auto_archiver/modules/hash_enricher/__manifest__.py index c7a023e..4f638de 100644 --- a/src/auto_archiver/modules/hash_enricher/__manifest__.py +++ b/src/auto_archiver/modules/hash_enricher/__manifest__.py @@ -3,16 +3,17 @@ "type": ["enricher"], "requires_setup": False, "dependencies": { - "python": ["loguru"], + "python": ["loguru"], }, "configs": { - "algorithm": {"default": "SHA-256", "help": "hash algorithm to use", "choices": ["SHA-256", "SHA3-512"]}, - # TODO add non-negative requirement to match previous implementation? - "chunksize": {"default": 16000000, - "help": "number of bytes to use when reading files in chunks (if this value is too large you will run out of RAM), default is 16MB", - 'type': 'int', - }, + "algorithm": {"default": "SHA-256", "help": "hash algorithm to use", "choices": ["SHA-256", "SHA3-512"]}, + # TODO add non-negative requirement to match previous implementation? + "chunksize": { + "default": 16000000, + "help": "number of bytes to use when reading files in chunks (if this value is too large you will run out of RAM), default is 16MB", + "type": "int", }, + }, "description": """ Generates cryptographic hashes for media files to ensure data integrity and authenticity. diff --git a/src/auto_archiver/modules/hash_enricher/hash_enricher.py b/src/auto_archiver/modules/hash_enricher/hash_enricher.py index 7a0587c..71425f2 100644 --- a/src/auto_archiver/modules/hash_enricher/hash_enricher.py +++ b/src/auto_archiver/modules/hash_enricher/hash_enricher.py @@ -1,4 +1,4 @@ -""" Hash Enricher for generating cryptographic hashes of media files. +"""Hash Enricher for generating cryptographic hashes of media files. The `HashEnricher` calculates cryptographic hashes (e.g., SHA-256, SHA3-512) for media files stored in `Metadata` objects. These hashes are used for @@ -7,6 +7,7 @@ exact duplicates. The hash is computed by reading the file's bytes in chunks, making it suitable for handling large files efficiently. """ + import hashlib from loguru import logger @@ -20,7 +21,6 @@ class HashEnricher(Enricher): Calculates hashes for Media instances """ - def enrich(self, to_enrich: Metadata) -> None: url = to_enrich.get_url() logger.debug(f"calculating media hashes for {url=} (using {self.algorithm})") @@ -35,5 +35,6 @@ class HashEnricher(Enricher): hash_algo = hashlib.sha256 elif self.algorithm == "SHA3-512": hash_algo = hashlib.sha3_512 - else: return "" + else: + return "" return calculate_file_hash(filename, hash_algo, self.chunksize) diff --git a/src/auto_archiver/modules/html_formatter/__init__.py b/src/auto_archiver/modules/html_formatter/__init__.py index 432ef33..fd1bb70 100644 --- a/src/auto_archiver/modules/html_formatter/__init__.py +++ b/src/auto_archiver/modules/html_formatter/__init__.py @@ -1 +1 @@ -from .html_formatter import HtmlFormatter \ No newline at end of file +from .html_formatter import HtmlFormatter diff --git a/src/auto_archiver/modules/html_formatter/__manifest__.py b/src/auto_archiver/modules/html_formatter/__manifest__.py index 6e51c7a..6501e4f 100644 --- a/src/auto_archiver/modules/html_formatter/__manifest__.py +++ b/src/auto_archiver/modules/html_formatter/__manifest__.py @@ -2,14 +2,13 @@ "name": "HTML Formatter", "type": ["formatter"], "requires_setup": False, - "dependencies": { - "python": ["hash_enricher", "loguru", "jinja2"], - "bin": [""] - }, + "dependencies": {"python": ["hash_enricher", "loguru", "jinja2"], "bin": [""]}, "configs": { - "detect_thumbnails": {"default": True, - "help": "if true will group by thumbnails generated by thumbnail enricher by id 'thumbnail_00'", - "type": "bool"}, + "detect_thumbnails": { + "default": True, + "help": "if true will group by thumbnails generated by thumbnail enricher by id 'thumbnail_00'", + "type": "bool", }, + }, "description": """ """, } diff --git a/src/auto_archiver/modules/html_formatter/html_formatter.py b/src/auto_archiver/modules/html_formatter/html_formatter.py index deb4b44..f5da1d8 100644 --- a/src/auto_archiver/modules/html_formatter/html_formatter.py +++ b/src/auto_archiver/modules/html_formatter/html_formatter.py @@ -1,5 +1,7 @@ from __future__ import annotations -import mimetypes, os, pathlib +import mimetypes +import os +import pathlib from jinja2 import Environment, FileSystemLoader from urllib.parse import quote from loguru import logger @@ -11,6 +13,7 @@ from auto_archiver.core import Metadata, Media from auto_archiver.core import Formatter from auto_archiver.utils.misc import random_str + class HtmlFormatter(Formatter): environment: Environment = None template: any = None @@ -21,9 +24,9 @@ class HtmlFormatter(Formatter): self.environment = Environment(loader=FileSystemLoader(template_dir), autoescape=True) # JinjaHelper class static methods are added as filters - self.environment.filters.update({ - k: v.__func__ for k, v in JinjaHelpers.__dict__.items() if isinstance(v, staticmethod) - }) + self.environment.filters.update( + {k: v.__func__ for k, v in JinjaHelpers.__dict__.items() if isinstance(v, staticmethod)} + ) # Load a specific template or default to "html_template.html" template_name = self.config.get("template_name", "html_template.html") @@ -36,11 +39,7 @@ class HtmlFormatter(Formatter): return content = self.template.render( - url=url, - title=item.get_title(), - media=item.media, - metadata=item.metadata, - version=__version__ + url=url, title=item.get_title(), media=item.media, metadata=item.metadata, version=__version__ ) html_path = os.path.join(self.tmp_dir, f"formatted{random_str(24)}.html") @@ -49,7 +48,7 @@ class HtmlFormatter(Formatter): final_media = Media(filename=html_path, _mimetype="text/html") # get the already instantiated hash_enricher module - he = self.module_factory.get_module('hash_enricher', self.config) + he = self.module_factory.get_module("hash_enricher", self.config) if len(hd := he.calculate_hash(final_media.filename)): final_media.set("hash", f"{he.algorithm}:{hd}") diff --git a/src/auto_archiver/modules/instagram_api_extractor/__manifest__.py b/src/auto_archiver/modules/instagram_api_extractor/__manifest__.py index 2d8f1d9..e10bd1e 100644 --- a/src/auto_archiver/modules/instagram_api_extractor/__manifest__.py +++ b/src/auto_archiver/modules/instagram_api_extractor/__manifest__.py @@ -2,18 +2,18 @@ "name": "Instagram API Extractor", "type": ["extractor"], "entry_point": "instagram_api_extractor::InstagramAPIExtractor", - "dependencies": - {"python": ["requests", - "loguru", - "retrying", - "tqdm",], - }, + "dependencies": { + "python": [ + "requests", + "loguru", + "retrying", + "tqdm", + ], + }, "requires_setup": True, "configs": { - "access_token": {"default": None, - "help": "a valid instagrapi-api token"}, - "api_endpoint": {"required": True, - "help": "API endpoint to use"}, + "access_token": {"default": None, "help": "a valid instagrapi-api token"}, + "api_endpoint": {"required": True, "help": "API endpoint to use"}, "full_profile": { "default": False, "type": "bool", diff --git a/src/auto_archiver/modules/instagram_api_extractor/instagram_api_extractor.py b/src/auto_archiver/modules/instagram_api_extractor/instagram_api_extractor.py index a75e065..5f13ecf 100644 --- a/src/auto_archiver/modules/instagram_api_extractor/instagram_api_extractor.py +++ b/src/auto_archiver/modules/instagram_api_extractor/instagram_api_extractor.py @@ -36,21 +36,16 @@ class InstagramAPIExtractor(Extractor): if self.api_endpoint[-1] == "/": self.api_endpoint = self.api_endpoint[:-1] - def download(self, item: Metadata) -> Metadata: url = item.get_url() - url.replace("instagr.com", "instagram.com").replace( - "instagr.am", "instagram.com" - ) + url.replace("instagr.com", "instagram.com").replace("instagr.am", "instagram.com") insta_matches = self.valid_url.findall(url) logger.info(f"{insta_matches=}") if not len(insta_matches) or len(insta_matches[0]) != 3: return if len(insta_matches) > 1: - logger.warning( - f"Multiple instagram matches found in {url=}, using the first one" - ) + logger.warning(f"Multiple instagram matches found in {url=}, using the first one") return g1, g2, g3 = insta_matches[0][0], insta_matches[0][1], insta_matches[0][2] if g1 == "": @@ -73,23 +68,20 @@ class InstagramAPIExtractor(Extractor): def call_api(self, path: str, params: dict) -> dict: headers = {"accept": "application/json", "x-access-key": self.access_token} logger.debug(f"calling {self.api_endpoint}/{path} with {params=}") - return requests.get( - f"{self.api_endpoint}/{path}", headers=headers, params=params - ).json() + return requests.get(f"{self.api_endpoint}/{path}", headers=headers, params=params).json() def cleanup_dict(self, d: dict | list) -> dict: # repeats 3 times to remove nested empty values if not self.minimize_json_output: return d - if type(d) == list: + if isinstance(d, list): return [self.cleanup_dict(v) for v in d] - if type(d) != dict: + if not isinstance(d, dict): return d return { k: clean_v for k, v in d.items() - if (clean_v := self.cleanup_dict(v)) - not in [0.0, 0, [], {}, "", None, "null"] + if (clean_v := self.cleanup_dict(v)) not in [0.0, 0, [], {}, "", None, "null"] and k not in ["x", "y", "width", "height"] } @@ -103,7 +95,7 @@ class InstagramAPIExtractor(Extractor): result.set_title(user.get("full_name", username)).set("data", user) if pic_url := user.get("profile_pic_url_hd", user.get("profile_pic_url")): filename = self.download_from_url(pic_url) - result.add_media(Media(filename=filename), id=f"profile_picture") + result.add_media(Media(filename=filename), id="profile_picture") if self.full_profile: user_id = user.get("pk") @@ -126,9 +118,7 @@ class InstagramAPIExtractor(Extractor): try: self.download_all_tagged(result, user_id) except Exception as e: - result.append( - "errors", f"Error downloading tagged posts for {username}" - ) + result.append("errors", f"Error downloading tagged posts for {username}") logger.error(f"Error downloading tagged posts for {username}: {e}") # download all highlights @@ -143,7 +133,7 @@ class InstagramAPIExtractor(Extractor): def download_all_highlights(self, result, username, user_id): count_highlights = 0 - highlights = self.call_api(f"v1/user/highlights", {"user_id": user_id}) + highlights = self.call_api("v1/user/highlights", {"user_id": user_id}) for h in highlights: try: h_info = self._download_highlights_reusable(result, h.get("pk")) @@ -153,26 +143,17 @@ class InstagramAPIExtractor(Extractor): "errors", f"Error downloading highlight id{h.get('pk')} for {username}", ) - logger.error( - f"Error downloading highlight id{h.get('pk')} for {username}: {e}" - ) - if ( - self.full_profile_max_posts - and count_highlights >= self.full_profile_max_posts - ): - logger.info( - f"HIGHLIGHTS reached full_profile_max_posts={self.full_profile_max_posts}" - ) + logger.error(f"Error downloading highlight id{h.get('pk')} for {username}: {e}") + if self.full_profile_max_posts and count_highlights >= self.full_profile_max_posts: + logger.info(f"HIGHLIGHTS reached full_profile_max_posts={self.full_profile_max_posts}") break result.set("#highlights", count_highlights) - def download_post( - self, result: Metadata, code: str = None, id: str = None, context: str = None - ) -> Metadata: + def download_post(self, result: Metadata, code: str = None, id: str = None, context: str = None) -> Metadata: if id: - post = self.call_api(f"v1/media/by/id", {"id": id}) + post = self.call_api("v1/media/by/id", {"id": id}) else: - post = self.call_api(f"v1/media/by/code", {"code": code}) + post = self.call_api("v1/media/by/code", {"code": code}) assert post, f"Post {id or code} not found" if caption_text := post.get("caption_text"): @@ -192,15 +173,11 @@ class InstagramAPIExtractor(Extractor): return result.success("insta highlights") def _download_highlights_reusable(self, result: Metadata, id: str) -> dict: - full_h = self.call_api(f"v2/highlight/by/id", {"id": id}) + full_h = self.call_api("v2/highlight/by/id", {"id": id}) h_info = full_h.get("response", {}).get("reels", {}).get(f"highlight:{id}") assert h_info, f"Highlight {id} not found: {full_h=}" - if ( - cover_media := h_info.get("cover_media", {}) - .get("cropped_image_version", {}) - .get("url") - ): + if cover_media := h_info.get("cover_media", {}).get("cropped_image_version", {}).get("url"): filename = self.download_from_url(cover_media) result.add_media(Media(filename=filename), id=f"cover_media highlight {id}") @@ -210,9 +187,7 @@ class InstagramAPIExtractor(Extractor): self.scrape_item(result, h, "highlight") except Exception as e: result.append("errors", f"Error downloading highlight {h.get('id')}") - logger.error( - f"Error downloading highlight, skipping {h.get('id')}: {e}" - ) + logger.error(f"Error downloading highlight, skipping {h.get('id')}: {e}") return h_info @@ -225,7 +200,7 @@ class InstagramAPIExtractor(Extractor): return result.success(f"insta stories {now}") def _download_stories_reusable(self, result: Metadata, username: str) -> list[dict]: - stories = self.call_api(f"v1/user/stories/by/username", {"username": username}) + stories = self.call_api("v1/user/stories/by/username", {"username": username}) if not stories or not len(stories): return [] stories = stories[::-1] # newest to oldest @@ -244,10 +219,8 @@ class InstagramAPIExtractor(Extractor): post_count = 0 while end_cursor != "": - posts = self.call_api( - f"v1/user/medias/chunk", {"user_id": user_id, "end_cursor": end_cursor} - ) - if not len(posts) or not type(posts) == list or len(posts) != 2: + posts = self.call_api("v1/user/medias/chunk", {"user_id": user_id, "end_cursor": end_cursor}) + if not posts or not isinstance(posts, list) or len(posts) != 2: break posts, end_cursor = posts[0], posts[1] logger.info(f"parsing {len(posts)} posts, next {end_cursor=}") @@ -260,13 +233,8 @@ class InstagramAPIExtractor(Extractor): logger.error(f"Error downloading post, skipping {p.get('id')}: {e}") pbar.update(1) post_count += 1 - if ( - self.full_profile_max_posts - and post_count >= self.full_profile_max_posts - ): - logger.info( - f"POSTS reached full_profile_max_posts={self.full_profile_max_posts}" - ) + if self.full_profile_max_posts and post_count >= self.full_profile_max_posts: + logger.info(f"POSTS reached full_profile_max_posts={self.full_profile_max_posts}") break result.set("#posts", post_count) @@ -275,10 +243,8 @@ class InstagramAPIExtractor(Extractor): pbar = tqdm(desc="downloading tagged posts") tagged_count = 0 - while next_page_id != None: - resp = self.call_api( - f"v2/user/tag/medias", {"user_id": user_id, "page_id": next_page_id} - ) + while next_page_id is not None: + resp = self.call_api("v2/user/tag/medias", {"user_id": user_id, "page_id": next_page_id}) posts = resp.get("response", {}).get("items", []) if not len(posts): break @@ -290,21 +256,12 @@ class InstagramAPIExtractor(Extractor): try: self.scrape_item(result, p, "tagged") except Exception as e: - result.append( - "errors", f"Error downloading tagged post {p.get('id')}" - ) - logger.error( - f"Error downloading tagged post, skipping {p.get('id')}: {e}" - ) + result.append("errors", f"Error downloading tagged post {p.get('id')}") + logger.error(f"Error downloading tagged post, skipping {p.get('id')}: {e}") pbar.update(1) tagged_count += 1 - if ( - self.full_profile_max_posts - and tagged_count >= self.full_profile_max_posts - ): - logger.info( - f"TAGS reached full_profile_max_posts={self.full_profile_max_posts}" - ) + if self.full_profile_max_posts and tagged_count >= self.full_profile_max_posts: + logger.info(f"TAGS reached full_profile_max_posts={self.full_profile_max_posts}") break result.set("#tagged", tagged_count) @@ -318,9 +275,7 @@ class InstagramAPIExtractor(Extractor): context can be used to give specific id prefixes to media """ if "clips_metadata" in item: - if reusable_text := item.get("clips_metadata", {}).get( - "reusable_text_attribute_string" - ): + if reusable_text := item.get("clips_metadata", {}).get("reusable_text_attribute_string"): item["clips_metadata_text"] = reusable_text if self.minimize_json_output: del item["clips_metadata"] diff --git a/src/auto_archiver/modules/instagram_extractor/__init__.py b/src/auto_archiver/modules/instagram_extractor/__init__.py index 6f39171..cefbcc5 100644 --- a/src/auto_archiver/modules/instagram_extractor/__init__.py +++ b/src/auto_archiver/modules/instagram_extractor/__init__.py @@ -1 +1 @@ -from .instagram_extractor import InstagramExtractor \ No newline at end of file +from .instagram_extractor import InstagramExtractor diff --git a/src/auto_archiver/modules/instagram_extractor/__manifest__.py b/src/auto_archiver/modules/instagram_extractor/__manifest__.py index a66389f..6067c92 100644 --- a/src/auto_archiver/modules/instagram_extractor/__manifest__.py +++ b/src/auto_archiver/modules/instagram_extractor/__manifest__.py @@ -9,8 +9,7 @@ }, "requires_setup": True, "configs": { - "username": {"required": True, - "help": "A valid Instagram username."}, + "username": {"required": True, "help": "A valid Instagram username."}, "password": { "required": True, "help": "The corresponding Instagram account password.", diff --git a/src/auto_archiver/modules/instagram_extractor/instagram_extractor.py b/src/auto_archiver/modules/instagram_extractor/instagram_extractor.py index 7e195ad..294b4e7 100644 --- a/src/auto_archiver/modules/instagram_extractor/instagram_extractor.py +++ b/src/auto_archiver/modules/instagram_extractor/instagram_extractor.py @@ -1,9 +1,12 @@ -""" Uses the Instaloader library to download content from Instagram. This class handles both individual posts - and user profiles, downloading as much information as possible, including images, videos, text, stories, - highlights, and tagged posts. Authentication is required via username/password or a session file. +"""Uses the Instaloader library to download content from Instagram. This class handles both individual posts +and user profiles, downloading as much information as possible, including images, videos, text, stories, +highlights, and tagged posts. Authentication is required via username/password or a session file. """ -import re, os, shutil + +import re +import os +import shutil import instaloader from loguru import logger @@ -11,6 +14,7 @@ from auto_archiver.core import Extractor from auto_archiver.core import Metadata from auto_archiver.core import Media + class InstagramExtractor(Extractor): """ Uses Instaloader to download either a post (inc images, videos, text) or as much as possible from a profile (posts, stories, highlights, ...) @@ -25,26 +29,24 @@ class InstagramExtractor(Extractor): # TODO: links to stories def setup(self) -> None: - self.insta = instaloader.Instaloader( download_geotags=True, download_comments=True, compress_json=False, dirname_pattern=self.download_folder, - filename_pattern="{date_utc}_UTC_{target}__{typename}" + filename_pattern="{date_utc}_UTC_{target}__{typename}", ) try: self.insta.load_session_from_file(self.username, self.session_file) - except Exception as e: + except Exception: try: - logger.debug(f"Session file failed", exc_info=True) + logger.debug("Session file failed", exc_info=True) logger.info("No valid session file found - Attempting login with use and password.") self.insta.login(self.username, self.password) self.insta.save_session_to_file(self.session_file) except Exception as e: logger.error(f"Failed to setup Instagram Extractor with Instagrapi. {e}") - def download(self, item: Metadata) -> Metadata: url = item.get_url() @@ -53,7 +55,8 @@ class InstagramExtractor(Extractor): profile_matches = self.profile_pattern.findall(url) # return if not a valid instagram link - if not len(post_matches) and not len(profile_matches): return + if not len(post_matches) and not len(profile_matches): + return result = None try: @@ -65,7 +68,9 @@ class InstagramExtractor(Extractor): elif len(profile_matches): result = self.download_profile(url, profile_matches[0]) except Exception as e: - logger.error(f"Failed to download with instagram extractor due to: {e}, make sure your account credentials are valid.") + logger.error( + f"Failed to download with instagram extractor due to: {e}, make sure your account credentials are valid." + ) finally: shutil.rmtree(self.download_folder, ignore_errors=True) return result @@ -84,35 +89,50 @@ class InstagramExtractor(Extractor): profile = instaloader.Profile.from_username(self.insta.context, username) try: for post in profile.get_posts(): - try: self.insta.download_post(post, target=f"profile_post_{post.owner_username}") - except Exception as e: logger.error(f"Failed to download post: {post.shortcode}: {e}") - except Exception as e: logger.error(f"Failed profile.get_posts: {e}") + try: + self.insta.download_post(post, target=f"profile_post_{post.owner_username}") + except Exception as e: + logger.error(f"Failed to download post: {post.shortcode}: {e}") + except Exception as e: + logger.error(f"Failed profile.get_posts: {e}") try: for post in profile.get_tagged_posts(): - try: self.insta.download_post(post, target=f"tagged_post_{post.owner_username}") - except Exception as e: logger.error(f"Failed to download tagged post: {post.shortcode}: {e}") - except Exception as e: logger.error(f"Failed profile.get_tagged_posts: {e}") + try: + self.insta.download_post(post, target=f"tagged_post_{post.owner_username}") + except Exception as e: + logger.error(f"Failed to download tagged post: {post.shortcode}: {e}") + except Exception as e: + logger.error(f"Failed profile.get_tagged_posts: {e}") try: for post in profile.get_igtv_posts(): - try: self.insta.download_post(post, target=f"igtv_post_{post.owner_username}") - except Exception as e: logger.error(f"Failed to download igtv post: {post.shortcode}: {e}") - except Exception as e: logger.error(f"Failed profile.get_igtv_posts: {e}") + try: + self.insta.download_post(post, target=f"igtv_post_{post.owner_username}") + except Exception as e: + logger.error(f"Failed to download igtv post: {post.shortcode}: {e}") + except Exception as e: + logger.error(f"Failed profile.get_igtv_posts: {e}") try: for story in self.insta.get_stories([profile.userid]): for item in story.get_items(): - try: self.insta.download_storyitem(item, target=f"story_item_{story.owner_username}") - except Exception as e: logger.error(f"Failed to download story item: {item}: {e}") - except Exception as e: logger.error(f"Failed get_stories: {e}") + try: + self.insta.download_storyitem(item, target=f"story_item_{story.owner_username}") + except Exception as e: + logger.error(f"Failed to download story item: {item}: {e}") + except Exception as e: + logger.error(f"Failed get_stories: {e}") try: for highlight in self.insta.get_highlights(profile.userid): for item in highlight.get_items(): - try: self.insta.download_storyitem(item, target=f"highlight_item_{highlight.owner_username}") - except Exception as e: logger.error(f"Failed to download highlight item: {item}: {e}") - except Exception as e: logger.error(f"Failed get_highlights: {e}") + try: + self.insta.download_storyitem(item, target=f"highlight_item_{highlight.owner_username}") + except Exception as e: + logger.error(f"Failed to download highlight item: {item}: {e}") + except Exception as e: + logger.error(f"Failed get_highlights: {e}") return self.process_downloads(url, f"@{username}", profile._asdict(), None) @@ -124,7 +144,8 @@ class InstagramExtractor(Extractor): all_media = [] for f in os.listdir(self.download_folder): if os.path.isfile((filename := os.path.join(self.download_folder, f))): - if filename[-4:] == ".txt": continue + if filename[-4:] == ".txt": + continue all_media.append(Media(filename)) assert len(all_media) > 1, "No uploaded media found" diff --git a/src/auto_archiver/modules/instagram_tbot_extractor/__manifest__.py b/src/auto_archiver/modules/instagram_tbot_extractor/__manifest__.py index a24a864..e3f94d0 100644 --- a/src/auto_archiver/modules/instagram_tbot_extractor/__manifest__.py +++ b/src/auto_archiver/modules/instagram_tbot_extractor/__manifest__.py @@ -1,16 +1,21 @@ { "name": "Instagram Telegram Bot Extractor", "type": ["extractor"], - "dependencies": {"python": ["loguru", "telethon",], - }, + "dependencies": { + "python": [ + "loguru", + "telethon", + ], + }, "requires_setup": True, "configs": { - "api_id": {"default": None, "help": "telegram API_ID value, go to https://my.telegram.org/apps"}, - "api_hash": {"default": None, "help": "telegram API_HASH value, go to https://my.telegram.org/apps"}, - "session_file": {"default": "secrets/anon-insta", "help": "optional, records the telegram login session for future usage, '.session' will be appended to the provided value."}, - "timeout": {"default": 45, - "type": "int", - "help": "timeout to fetch the instagram content in seconds."}, + "api_id": {"default": None, "help": "telegram API_ID value, go to https://my.telegram.org/apps"}, + "api_hash": {"default": None, "help": "telegram API_HASH value, go to https://my.telegram.org/apps"}, + "session_file": { + "default": "secrets/anon-insta", + "help": "optional, records the telegram login session for future usage, '.session' will be appended to the provided value.", + }, + "timeout": {"default": 45, "type": "int", "help": "timeout to fetch the instagram content in seconds."}, }, "description": """ The `InstagramTbotExtractor` module uses a Telegram bot (`instagram_load_bot`) to fetch and archive Instagram content, diff --git a/src/auto_archiver/modules/instagram_tbot_extractor/instagram_tbot_extractor.py b/src/auto_archiver/modules/instagram_tbot_extractor/instagram_tbot_extractor.py index 1416da9..81d2bf6 100644 --- a/src/auto_archiver/modules/instagram_tbot_extractor/instagram_tbot_extractor.py +++ b/src/auto_archiver/modules/instagram_tbot_extractor/instagram_tbot_extractor.py @@ -51,7 +51,7 @@ class InstagramTbotExtractor(Extractor): """Initializes the Telegram client.""" try: self.client = TelegramClient(self.session_file, self.api_id, self.api_hash) - except OperationalError as e: + except OperationalError: logger.error( f"Unable to access the {self.session_file} session. " "Ensure that you don't use the same session file here and in telethon_extractor. " @@ -65,15 +65,15 @@ class InstagramTbotExtractor(Extractor): session_file_name = self.session_file + ".session" if os.path.exists(session_file_name): os.remove(session_file_name) - + def download(self, item: Metadata) -> Metadata: url = item.get_url() - if not "instagram.com" in url: return False + if "instagram.com" not in url: + return False result = Metadata() tmp_dir = self.tmp_dir with self.client.start(): - chat, since_id = self._send_url_to_bot(url) message = self._process_messages(chat, since_id, tmp_dir, result) @@ -110,13 +110,14 @@ class InstagramTbotExtractor(Extractor): for post in self.client.iter_messages(chat, min_id=since_id): since_id = max(since_id, post.id) # Skip known filler message: - if post.message == 'The bot receives information through https://hikerapi.com/p/hJqpppqi': + if post.message == "The bot receives information through https://hikerapi.com/p/hJqpppqi": continue if post.media and post.id not in seen_media: - filename_dest = os.path.join(tmp_dir, f'{chat.id}_{post.id}') + filename_dest = os.path.join(tmp_dir, f"{chat.id}_{post.id}") media = self.client.download_media(post.media, filename_dest) if media: result.add_media(Media(media)) seen_media.append(post.id) - if post.message: message += post.message - return message.strip() \ No newline at end of file + if post.message: + message += post.message + return message.strip() diff --git a/src/auto_archiver/modules/local_storage/__init__.py b/src/auto_archiver/modules/local_storage/__init__.py index d23147d..e9c81f8 100644 --- a/src/auto_archiver/modules/local_storage/__init__.py +++ b/src/auto_archiver/modules/local_storage/__init__.py @@ -1 +1 @@ -from .local_storage import LocalStorage \ No newline at end of file +from .local_storage import LocalStorage diff --git a/src/auto_archiver/modules/local_storage/__manifest__.py b/src/auto_archiver/modules/local_storage/__manifest__.py index 72f59d1..b5307b8 100644 --- a/src/auto_archiver/modules/local_storage/__manifest__.py +++ b/src/auto_archiver/modules/local_storage/__manifest__.py @@ -17,9 +17,11 @@ "choices": ["random", "static"], }, "save_to": {"default": "./local_archive", "help": "folder where to save archived content"}, - "save_absolute": {"default": False, - "type": "bool", - "help": "whether the path to the stored file is absolute or relative in the output result inc. formatters (WARN: leaks the file structure)"}, + "save_absolute": { + "default": False, + "type": "bool", + "help": "whether the path to the stored file is absolute or relative in the output result inc. formatters (WARN: leaks the file structure)", + }, }, "description": """ LocalStorage: A storage module for saving archived content locally on the filesystem. @@ -33,5 +35,5 @@ ### Notes - Default storage folder is `./archived`, but this can be changed via the `save_to` configuration. - The `save_absolute` option can reveal the file structure in output formats; use with caution. - """ + """, } diff --git a/src/auto_archiver/modules/local_storage/local_storage.py b/src/auto_archiver/modules/local_storage/local_storage.py index 2b1a101..fdc6978 100644 --- a/src/auto_archiver/modules/local_storage/local_storage.py +++ b/src/auto_archiver/modules/local_storage/local_storage.py @@ -1,4 +1,3 @@ - import shutil from typing import IO import os @@ -8,12 +7,13 @@ from auto_archiver.core import Media from auto_archiver.core import Storage from auto_archiver.core.consts import SetupError + class LocalStorage(Storage): - - def setup(self) -> None: if len(self.save_to) > 200: - raise SetupError(f"Your save_to path is too long, this will cause issues saving files on your computer. Please use a shorter path.") + raise SetupError( + "Your save_to path is too long, this will cause issues saving files on your computer. Please use a shorter path." + ) def get_cdn_url(self, media: Media) -> str: dest = media.key @@ -25,18 +25,18 @@ class LocalStorage(Storage): def set_key(self, media, url, metadata): # clarify we want to save the file to the save_to folder - old_folder = metadata.get('folder', '') - metadata.set_context('folder', os.path.join(self.save_to, metadata.get('folder', ''))) + old_folder = metadata.get("folder", "") + metadata.set_context("folder", os.path.join(self.save_to, metadata.get("folder", ""))) super().set_key(media, url, metadata) # don't impact other storages that might want a different 'folder' set - metadata.set_context('folder', old_folder) + metadata.set_context("folder", old_folder) def upload(self, media: Media, **kwargs) -> bool: # override parent so that we can use shutil.copy2 and keep metadata dest = media.key os.makedirs(os.path.dirname(dest), exist_ok=True) - logger.debug(f'[{self.__class__.__name__}] storing file {media.filename} with key {media.key} to {dest}') + logger.debug(f"[{self.__class__.__name__}] storing file {media.filename} with key {media.key} to {dest}") res = shutil.copy2(media.filename, dest) logger.info(res) @@ -44,4 +44,4 @@ class LocalStorage(Storage): # must be implemented even if unused def uploadf(self, file: IO[bytes], key: str, **kwargs: dict) -> bool: - pass \ No newline at end of file + pass diff --git a/src/auto_archiver/modules/meta_enricher/__manifest__.py b/src/auto_archiver/modules/meta_enricher/__manifest__.py index 37c9201..d3732a3 100644 --- a/src/auto_archiver/modules/meta_enricher/__manifest__.py +++ b/src/auto_archiver/modules/meta_enricher/__manifest__.py @@ -3,7 +3,7 @@ "type": ["enricher"], "requires_setup": False, "dependencies": { - "python": ["loguru"], + "python": ["loguru"], }, "description": """ Adds metadata information about the archive operations, Adds metadata about archive operations, including file sizes and archive duration./ diff --git a/src/auto_archiver/modules/meta_enricher/meta_enricher.py b/src/auto_archiver/modules/meta_enricher/meta_enricher.py index 03fb01e..9356b16 100644 --- a/src/auto_archiver/modules/meta_enricher/meta_enricher.py +++ b/src/auto_archiver/modules/meta_enricher/meta_enricher.py @@ -23,7 +23,9 @@ class MetaEnricher(Enricher): self.enrich_archive_duration(to_enrich) def enrich_file_sizes(self, to_enrich: Metadata): - logger.debug(f"calculating archive file sizes for url={to_enrich.get_url()} ({len(to_enrich.media)} media files)") + logger.debug( + f"calculating archive file sizes for url={to_enrich.get_url()} ({len(to_enrich.media)} media files)" + ) total_size = 0 for media in to_enrich.get_all_media(): file_stats = os.stat(media.filename) @@ -34,7 +36,6 @@ class MetaEnricher(Enricher): to_enrich.set("total_bytes", total_size) to_enrich.set("total_size", self.human_readable_bytes(total_size)) - def human_readable_bytes(self, size: int) -> str: # receives number of bytes and returns human readble size for unit in ["bytes", "KB", "MB", "GB", "TB"]: @@ -46,4 +47,4 @@ class MetaEnricher(Enricher): logger.debug(f"calculating archive duration for url={to_enrich.get_url()} ") archive_duration = datetime.datetime.now(datetime.timezone.utc) - to_enrich.get("_processed_at") - to_enrich.set("archive_duration_seconds", archive_duration.seconds) \ No newline at end of file + to_enrich.set("archive_duration_seconds", archive_duration.seconds) diff --git a/src/auto_archiver/modules/metadata_enricher/__init__.py b/src/auto_archiver/modules/metadata_enricher/__init__.py index 020bd4a..2fe894a 100644 --- a/src/auto_archiver/modules/metadata_enricher/__init__.py +++ b/src/auto_archiver/modules/metadata_enricher/__init__.py @@ -1 +1 @@ -from .metadata_enricher import MetadataEnricher \ No newline at end of file +from .metadata_enricher import MetadataEnricher diff --git a/src/auto_archiver/modules/metadata_enricher/__manifest__.py b/src/auto_archiver/modules/metadata_enricher/__manifest__.py index f8ccdc6..3727551 100644 --- a/src/auto_archiver/modules/metadata_enricher/__manifest__.py +++ b/src/auto_archiver/modules/metadata_enricher/__manifest__.py @@ -2,10 +2,7 @@ "name": "Media Metadata Enricher", "type": ["enricher"], "requires_setup": True, - "dependencies": { - "python": ["loguru"], - "bin": ["exiftool"] - }, + "dependencies": {"python": ["loguru"], "bin": ["exiftool"]}, "description": """ Extracts metadata information from files using ExifTool. @@ -17,5 +14,5 @@ ### Notes - Requires ExifTool to be installed and accessible via the system's PATH. - Skips enrichment for files where metadata extraction fails. - """ + """, } diff --git a/src/auto_archiver/modules/metadata_enricher/metadata_enricher.py b/src/auto_archiver/modules/metadata_enricher/metadata_enricher.py index c052d0a..e4fac44 100644 --- a/src/auto_archiver/modules/metadata_enricher/metadata_enricher.py +++ b/src/auto_archiver/modules/metadata_enricher/metadata_enricher.py @@ -11,7 +11,6 @@ class MetadataEnricher(Enricher): Extracts metadata information from files using exiftool. """ - def enrich(self, to_enrich: Metadata) -> None: url = to_enrich.get_url() logger.debug(f"extracting EXIF metadata for {url=}") @@ -23,13 +22,13 @@ class MetadataEnricher(Enricher): def get_metadata(self, filename: str) -> dict: try: # Run ExifTool command to extract metadata from the file - cmd = ['exiftool', filename] + cmd = ["exiftool", filename] result = subprocess.run(cmd, capture_output=True, text=True) # Process the output to extract individual metadata fields metadata = {} for line in result.stdout.splitlines(): - field, value = line.strip().split(':', 1) + field, value = line.strip().split(":", 1) metadata[field.strip()] = value.strip() return metadata except FileNotFoundError: diff --git a/src/auto_archiver/modules/mute_formatter/__manifest__.py b/src/auto_archiver/modules/mute_formatter/__manifest__.py index e81dc4c..185645e 100644 --- a/src/auto_archiver/modules/mute_formatter/__manifest__.py +++ b/src/auto_archiver/modules/mute_formatter/__manifest__.py @@ -2,8 +2,7 @@ "name": "Mute Formatter", "type": ["formatter"], "requires_setup": True, - "dependencies": { - }, + "dependencies": {}, "description": """ Default formatter. """, } diff --git a/src/auto_archiver/modules/mute_formatter/mute_formatter.py b/src/auto_archiver/modules/mute_formatter/mute_formatter.py index 129ddcb..b7c0ba5 100644 --- a/src/auto_archiver/modules/mute_formatter/mute_formatter.py +++ b/src/auto_archiver/modules/mute_formatter/mute_formatter.py @@ -5,5 +5,5 @@ from auto_archiver.core import Formatter class MuteFormatter(Formatter): - - def format(self, item: Metadata) -> Media: return None + def format(self, item: Metadata) -> Media: + return None diff --git a/src/auto_archiver/modules/pdq_hash_enricher/__init__.py b/src/auto_archiver/modules/pdq_hash_enricher/__init__.py index b444197..88a964b 100644 --- a/src/auto_archiver/modules/pdq_hash_enricher/__init__.py +++ b/src/auto_archiver/modules/pdq_hash_enricher/__init__.py @@ -1 +1 @@ -from .pdq_hash_enricher import PdqHashEnricher \ No newline at end of file +from .pdq_hash_enricher import PdqHashEnricher diff --git a/src/auto_archiver/modules/pdq_hash_enricher/__manifest__.py b/src/auto_archiver/modules/pdq_hash_enricher/__manifest__.py index 133fef7..9c7a5c8 100644 --- a/src/auto_archiver/modules/pdq_hash_enricher/__manifest__.py +++ b/src/auto_archiver/modules/pdq_hash_enricher/__manifest__.py @@ -17,5 +17,5 @@ ### Notes - Best used after enrichers like `thumbnail_enricher` or `screenshot_enricher` to ensure images are available. - Uses the `pdqhash` library to compute 256-bit perceptual hashes, which are stored as hexadecimal strings. - """ + """, } diff --git a/src/auto_archiver/modules/pdq_hash_enricher/pdq_hash_enricher.py b/src/auto_archiver/modules/pdq_hash_enricher/pdq_hash_enricher.py index e812e8b..c7d4a47 100644 --- a/src/auto_archiver/modules/pdq_hash_enricher/pdq_hash_enricher.py +++ b/src/auto_archiver/modules/pdq_hash_enricher/pdq_hash_enricher.py @@ -10,6 +10,7 @@ This enricher is typically used after thumbnail or screenshot enrichers to ensure images are available for hashing. """ + import traceback import pdqhash import numpy as np @@ -34,7 +35,12 @@ class PdqHashEnricher(Enricher): for m in to_enrich.media: for media in m.all_inner_media(True): media_id = media.get("id", "") - if media.is_image() and "screenshot" not in media_id and "warc-file-" not in media_id and len(hd := self.calculate_pdq_hash(media.filename)): + if ( + media.is_image() + and "screenshot" not in media_id + and "warc-file-" not in media_id + and len(hd := self.calculate_pdq_hash(media.filename)) + ): media.set("pdq_hash", hd) media_with_hashes.append(media.filename) @@ -51,5 +57,7 @@ class PdqHashEnricher(Enricher): hash = "".join(str(b) for b in hash_array) return hex(int(hash, 2))[2:] except UnidentifiedImageError as e: - logger.error(f"Image {filename=} is likely corrupted or in unsupported format {e}: {traceback.format_exc()}") + logger.error( + f"Image {filename=} is likely corrupted or in unsupported format {e}: {traceback.format_exc()}" + ) return "" diff --git a/src/auto_archiver/modules/s3_storage/__init__.py b/src/auto_archiver/modules/s3_storage/__init__.py index cbf3237..5e388d1 100644 --- a/src/auto_archiver/modules/s3_storage/__init__.py +++ b/src/auto_archiver/modules/s3_storage/__init__.py @@ -1 +1 @@ -from .s3_storage import S3Storage \ No newline at end of file +from .s3_storage import S3Storage diff --git a/src/auto_archiver/modules/s3_storage/__manifest__.py b/src/auto_archiver/modules/s3_storage/__manifest__.py index 156f562..3118d0e 100644 --- a/src/auto_archiver/modules/s3_storage/__manifest__.py +++ b/src/auto_archiver/modules/s3_storage/__manifest__.py @@ -20,20 +20,20 @@ "region": {"default": None, "help": "S3 region name"}, "key": {"default": None, "help": "S3 API key"}, "secret": {"default": None, "help": "S3 API secret"}, - "random_no_duplicate": {"default": False, - "type": "bool", - "help": "if set, it will override `path_generator`, `filename_generator` and `folder`. It will check if the file already exists and if so it will not upload it again. Creates a new root folder path `no-dups/`"}, + "random_no_duplicate": { + "default": False, + "type": "bool", + "help": "if set, it will override `path_generator`, `filename_generator` and `folder`. It will check if the file already exists and if so it will not upload it again. Creates a new root folder path `no-dups/`", + }, "endpoint_url": { - "default": 'https://{region}.digitaloceanspaces.com', - "help": "S3 bucket endpoint, {region} are inserted at runtime" + "default": "https://{region}.digitaloceanspaces.com", + "help": "S3 bucket endpoint, {region} are inserted at runtime", }, "cdn_url": { - "default": 'https://{bucket}.{region}.cdn.digitaloceanspaces.com/{key}', - "help": "S3 CDN url, {bucket}, {region} and {key} are inserted at runtime" + "default": "https://{bucket}.{region}.cdn.digitaloceanspaces.com/{key}", + "help": "S3 CDN url, {bucket}, {region} and {key} are inserted at runtime", }, - "private": {"default": False, - "type": "bool", - "help": "if true S3 files will not be readable online"}, + "private": {"default": False, "type": "bool", "help": "if true S3 files will not be readable online"}, }, "description": """ S3Storage: A storage module for saving media files to an S3-compatible object storage. @@ -50,5 +50,5 @@ - The `random_no_duplicate` option ensures no duplicate uploads by leveraging hash-based folder structures. - Uses `boto3` for interaction with the S3 API. - Depends on the `HashEnricher` module for hash calculation. - """ + """, } diff --git a/src/auto_archiver/modules/s3_storage/s3_storage.py b/src/auto_archiver/modules/s3_storage/s3_storage.py index 183a944..abac4f7 100644 --- a/src/auto_archiver/modules/s3_storage/s3_storage.py +++ b/src/auto_archiver/modules/s3_storage/s3_storage.py @@ -1,4 +1,3 @@ - from typing import IO import boto3 @@ -11,33 +10,36 @@ from auto_archiver.utils.misc import calculate_file_hash, random_str NO_DUPLICATES_FOLDER = "no-dups/" -class S3Storage(Storage): +class S3Storage(Storage): def setup(self) -> None: self.s3 = boto3.client( - 's3', + "s3", region_name=self.region, endpoint_url=self.endpoint_url.format(region=self.region), aws_access_key_id=self.key, - aws_secret_access_key=self.secret + aws_secret_access_key=self.secret, ) if self.random_no_duplicate: - logger.warning("random_no_duplicate is set to True, this will override `path_generator`, `filename_generator` and `folder`.") + logger.warning( + "random_no_duplicate is set to True, this will override `path_generator`, `filename_generator` and `folder`." + ) def get_cdn_url(self, media: Media) -> str: return self.cdn_url.format(bucket=self.bucket, region=self.region, key=media.key) def uploadf(self, file: IO[bytes], media: Media, **kwargs: dict) -> None: - if not self.is_upload_needed(media): return True + if not self.is_upload_needed(media): + return True extra_args = kwargs.get("extra_args", {}) - if not self.private and 'ACL' not in extra_args: - extra_args['ACL'] = 'public-read' + if not self.private and "ACL" not in extra_args: + extra_args["ACL"] = "public-read" - if 'ContentType' not in extra_args: + if "ContentType" not in extra_args: try: if media.mimetype: - extra_args['ContentType'] = media.mimetype + extra_args["ContentType"] = media.mimetype except Exception as e: logger.warning(f"Unable to get mimetype for {media.key=}, error: {e}") self.s3.upload_fileobj(file, Bucket=self.bucket, Key=media.key, ExtraArgs=extra_args) @@ -49,21 +51,21 @@ class S3Storage(Storage): hd = calculate_file_hash(media.filename) path = os.path.join(NO_DUPLICATES_FOLDER, hd[:24]) - if existing_key:=self.file_in_folder(path): + if existing_key := self.file_in_folder(path): media._key = existing_key media.set("previously archived", True) logger.debug(f"skipping upload of {media.filename} because it already exists in {media.key}") return False - + _, ext = os.path.splitext(media.key) media._key = os.path.join(path, f"{random_str(24)}{ext}") return True - def file_in_folder(self, path:str) -> str: + def file_in_folder(self, path: str) -> str: # checks if path exists and is not an empty folder - if not path.endswith('/'): - path = path + '/' - resp = self.s3.list_objects(Bucket=self.bucket, Prefix=path, Delimiter='/', MaxKeys=1) - if 'Contents' in resp: - return resp['Contents'][0]['Key'] - return False \ No newline at end of file + if not path.endswith("/"): + path = path + "/" + resp = self.s3.list_objects(Bucket=self.bucket, Prefix=path, Delimiter="/", MaxKeys=1) + if "Contents" in resp: + return resp["Contents"][0]["Key"] + return False diff --git a/src/auto_archiver/modules/screenshot_enricher/__manifest__.py b/src/auto_archiver/modules/screenshot_enricher/__manifest__.py index 92c0883..db04e6c 100644 --- a/src/auto_archiver/modules/screenshot_enricher/__manifest__.py +++ b/src/auto_archiver/modules/screenshot_enricher/__manifest__.py @@ -6,26 +6,29 @@ "python": ["loguru", "selenium"], }, "configs": { - "width": {"default": 1280, - "type": "int", - "help": "width of the screenshots"}, - "height": {"default": 1024, - "type": "int", - "help": "height of the screenshots"}, - "timeout": {"default": 60, - "type": "int", - "help": "timeout for taking the screenshot"}, - "sleep_before_screenshot": {"default": 4, - "type": "int", - "help": "seconds to wait for the pages to load before taking screenshot"}, - "http_proxy": {"default": "", "help": "http proxy to use for the webdriver, eg http://proxy-user:password@proxy-ip:port"}, - "save_to_pdf": {"default": False, - "type": "bool", - "help": "save the page as pdf along with the screenshot. PDF saving options can be adjusted with the 'print_options' parameter"}, - "print_options": {"default": {}, - "help": "options to pass to the pdf printer, in JSON format. See https://www.selenium.dev/documentation/webdriver/interactions/print_page/ for more information", - "type": "json_loader"}, + "width": {"default": 1280, "type": "int", "help": "width of the screenshots"}, + "height": {"default": 1024, "type": "int", "help": "height of the screenshots"}, + "timeout": {"default": 60, "type": "int", "help": "timeout for taking the screenshot"}, + "sleep_before_screenshot": { + "default": 4, + "type": "int", + "help": "seconds to wait for the pages to load before taking screenshot", }, + "http_proxy": { + "default": "", + "help": "http proxy to use for the webdriver, eg http://proxy-user:password@proxy-ip:port", + }, + "save_to_pdf": { + "default": False, + "type": "bool", + "help": "save the page as pdf along with the screenshot. PDF saving options can be adjusted with the 'print_options' parameter", + }, + "print_options": { + "default": {}, + "help": "options to pass to the pdf printer, in JSON format. See https://www.selenium.dev/documentation/webdriver/interactions/print_page/ for more information", + "type": "json_loader", + }, + }, "description": """ Captures screenshots and optionally saves web pages as PDFs using a WebDriver. @@ -37,5 +40,5 @@ ### Notes - Requires a WebDriver (e.g., ChromeDriver) installed and accessible via the system's PATH. - """ + """, } diff --git a/src/auto_archiver/modules/screenshot_enricher/screenshot_enricher.py b/src/auto_archiver/modules/screenshot_enricher/screenshot_enricher.py index 832d0f8..491bd51 100644 --- a/src/auto_archiver/modules/screenshot_enricher/screenshot_enricher.py +++ b/src/auto_archiver/modules/screenshot_enricher/screenshot_enricher.py @@ -1,5 +1,6 @@ from loguru import logger -import time, os +import time +import os import base64 from selenium.common.exceptions import TimeoutException @@ -9,8 +10,8 @@ from auto_archiver.core import Enricher from auto_archiver.utils import Webdriver, url as UrlUtil, random_str from auto_archiver.core import Media, Metadata -class ScreenshotEnricher(Enricher): +class ScreenshotEnricher(Enricher): def __init__(self, webdriver_factory=None): super().__init__() self.webdriver_factory = webdriver_factory or Webdriver @@ -25,8 +26,14 @@ class ScreenshotEnricher(Enricher): logger.debug(f"Enriching screenshot for {url=}") auth = self.auth_for_site(url) with self.webdriver_factory( - self.width, self.height, self.timeout, facebook_accept_cookies='facebook.com' in url, - http_proxy=self.http_proxy, print_options=self.print_options, auth=auth) as driver: + self.width, + self.height, + self.timeout, + facebook_accept_cookies="facebook.com" in url, + http_proxy=self.http_proxy, + print_options=self.print_options, + auth=auth, + ) as driver: try: driver.get(url) time.sleep(int(self.sleep_before_screenshot)) @@ -43,4 +50,3 @@ class ScreenshotEnricher(Enricher): logger.info("TimeoutException loading page for screenshot") except Exception as e: logger.error(f"Got error while loading webdriver for screenshot enricher: {e}") - diff --git a/src/auto_archiver/modules/ssl_enricher/__init__.py b/src/auto_archiver/modules/ssl_enricher/__init__.py index 23d2bee..86b9638 100644 --- a/src/auto_archiver/modules/ssl_enricher/__init__.py +++ b/src/auto_archiver/modules/ssl_enricher/__init__.py @@ -1 +1 @@ -from .ssl_enricher import SSLEnricher \ No newline at end of file +from .ssl_enricher import SSLEnricher diff --git a/src/auto_archiver/modules/ssl_enricher/__manifest__.py b/src/auto_archiver/modules/ssl_enricher/__manifest__.py index 097cd21..959fe2f 100644 --- a/src/auto_archiver/modules/ssl_enricher/__manifest__.py +++ b/src/auto_archiver/modules/ssl_enricher/__manifest__.py @@ -5,11 +5,13 @@ "dependencies": { "python": ["loguru", "slugify"], }, - 'entry_point': 'ssl_enricher::SSLEnricher', + "entry_point": "ssl_enricher::SSLEnricher", "configs": { - "skip_when_nothing_archived": {"default": True, - "type": 'bool', - "help": "if true, will skip enriching when no media is archived"}, + "skip_when_nothing_archived": { + "default": True, + "type": "bool", + "help": "if true, will skip enriching when no media is archived", + }, }, "description": """ Retrieves SSL certificate information for a domain and stores it as a file. @@ -21,5 +23,5 @@ ### Notes - Requires the target URL to use the HTTPS scheme; other schemes are not supported. - """ + """, } diff --git a/src/auto_archiver/modules/ssl_enricher/ssl_enricher.py b/src/auto_archiver/modules/ssl_enricher/ssl_enricher.py index b429163..3ab1389 100644 --- a/src/auto_archiver/modules/ssl_enricher/ssl_enricher.py +++ b/src/auto_archiver/modules/ssl_enricher/ssl_enricher.py @@ -1,4 +1,5 @@ -import ssl, os +import ssl +import os from slugify import slugify from urllib.parse import urlparse from loguru import logger @@ -13,16 +14,18 @@ class SSLEnricher(Enricher): """ def enrich(self, to_enrich: Metadata) -> None: - if not to_enrich.media and self.skip_when_nothing_archived: return - + if not to_enrich.media and self.skip_when_nothing_archived: + return + url = to_enrich.get_url() parsed = urlparse(url) assert parsed.scheme in ["https"], f"Invalid URL scheme {url=}" - + domain = parsed.netloc logger.debug(f"fetching SSL certificate for {domain=} in {url=}") cert = ssl.get_server_certificate((domain, 443)) cert_fn = os.path.join(self.tmp_dir, f"{slugify(domain)}.pem") - with open(cert_fn, "w") as f: f.write(cert) + with open(cert_fn, "w") as f: + f.write(cert) to_enrich.add_media(Media(filename=cert_fn), id="ssl_certificate") diff --git a/src/auto_archiver/modules/telegram_extractor/__init__.py b/src/auto_archiver/modules/telegram_extractor/__init__.py index 1fd80c2..18b73c2 100644 --- a/src/auto_archiver/modules/telegram_extractor/__init__.py +++ b/src/auto_archiver/modules/telegram_extractor/__init__.py @@ -1 +1 @@ -from .telegram_extractor import TelegramExtractor \ No newline at end of file +from .telegram_extractor import TelegramExtractor diff --git a/src/auto_archiver/modules/telegram_extractor/telegram_extractor.py b/src/auto_archiver/modules/telegram_extractor/telegram_extractor.py index d612e24..e70198d 100644 --- a/src/auto_archiver/modules/telegram_extractor/telegram_extractor.py +++ b/src/auto_archiver/modules/telegram_extractor/telegram_extractor.py @@ -1,4 +1,6 @@ -import requests, re, html +import requests +import re +import html from bs4 import BeautifulSoup from loguru import logger @@ -15,11 +17,11 @@ class TelegramExtractor(Extractor): def download(self, item: Metadata) -> Metadata: url = item.get_url() # detect URLs that we definitely cannot handle - if 't.me' != item.netloc: + if "t.me" != item.netloc: return False headers = { - 'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/81.0.4044.138 Safari/537.36' + "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/81.0.4044.138 Safari/537.36" } # TODO: check if we can do this more resilient to variable URLs @@ -27,11 +29,11 @@ class TelegramExtractor(Extractor): url += "?embed=1" t = requests.get(url, headers=headers) - s = BeautifulSoup(t.content, 'html.parser') + s = BeautifulSoup(t.content, "html.parser") result = Metadata() result.set_content(html.escape(str(t.content))) - if (timestamp := (s.find_all('time') or [{}])[0].get('datetime')): + if timestamp := (s.find_all("time") or [{}])[0].get("datetime"): result.set_timestamp(timestamp) video = s.find("video") @@ -41,25 +43,26 @@ class TelegramExtractor(Extractor): image_urls = [] for im in image_tags: - urls = [u.replace("'", "") for u in re.findall(r'url\((.*?)\)', im['style'])] + urls = [u.replace("'", "") for u in re.findall(r"url\((.*?)\)", im["style"])] image_urls += urls - if not len(image_urls): return False + if not len(image_urls): + return False for img_url in image_urls: result.add_media(Media(self.download_from_url(img_url))) else: - video_url = video.get('src') + video_url = video.get("src") m_video = Media(self.download_from_url(video_url)) # extract duration from HTML try: - duration = s.find_all('time')[0].contents[0] - if ':' in duration: - duration = float(duration.split( - ':')[0]) * 60 + float(duration.split(':')[1]) + duration = s.find_all("time")[0].contents[0] + if ":" in duration: + duration = float(duration.split(":")[0]) * 60 + float(duration.split(":")[1]) else: duration = float(duration) m_video.set("duration", duration) - except: pass + except Exception: + pass result.add_media(m_video) return result.success("telegram") diff --git a/src/auto_archiver/modules/telethon_extractor/__init__.py b/src/auto_archiver/modules/telethon_extractor/__init__.py index 2eaa57c..9d5e963 100644 --- a/src/auto_archiver/modules/telethon_extractor/__init__.py +++ b/src/auto_archiver/modules/telethon_extractor/__init__.py @@ -1 +1 @@ -from .telethon_extractor import TelethonExtractor \ No newline at end of file +from .telethon_extractor import TelethonExtractor diff --git a/src/auto_archiver/modules/telethon_extractor/__manifest__.py b/src/auto_archiver/modules/telethon_extractor/__manifest__.py index 5e58203..150b62c 100644 --- a/src/auto_archiver/modules/telethon_extractor/__manifest__.py +++ b/src/auto_archiver/modules/telethon_extractor/__manifest__.py @@ -3,26 +3,35 @@ "type": ["extractor"], "requires_setup": True, "dependencies": { - "python": ["telethon", - "loguru", - "tqdm", - ], - "bin": [""] + "python": [ + "telethon", + "loguru", + "tqdm", + ], + "bin": [""], }, "configs": { - "api_id": {"default": None, "help": "telegram API_ID value, go to https://my.telegram.org/apps"}, - "api_hash": {"default": None, "help": "telegram API_HASH value, go to https://my.telegram.org/apps"}, - "bot_token": {"default": None, "help": "optional, but allows access to more content such as large videos, talk to @botfather"}, - "session_file": {"default": "secrets/anon", "help": "optional, records the telegram login session for future usage, '.session' will be appended to the provided value."}, - "join_channels": {"default": True, - "type": "bool", - "help": "disables the initial setup with channel_invites config, useful if you have a lot and get stuck"}, - "channel_invites": { - "default": {}, - "help": "(JSON string) private channel invite links (format: t.me/joinchat/HASH OR t.me/+HASH) and (optional but important to avoid hanging for minutes on startup) channel id (format: CHANNEL_ID taken from a post url like https://t.me/c/CHANNEL_ID/1), the telegram account will join any new channels on setup", - "type": "json_loader", - } + "api_id": {"default": None, "help": "telegram API_ID value, go to https://my.telegram.org/apps"}, + "api_hash": {"default": None, "help": "telegram API_HASH value, go to https://my.telegram.org/apps"}, + "bot_token": { + "default": None, + "help": "optional, but allows access to more content such as large videos, talk to @botfather", }, + "session_file": { + "default": "secrets/anon", + "help": "optional, records the telegram login session for future usage, '.session' will be appended to the provided value.", + }, + "join_channels": { + "default": True, + "type": "bool", + "help": "disables the initial setup with channel_invites config, useful if you have a lot and get stuck", + }, + "channel_invites": { + "default": {}, + "help": "(JSON string) private channel invite links (format: t.me/joinchat/HASH OR t.me/+HASH) and (optional but important to avoid hanging for minutes on startup) channel id (format: CHANNEL_ID taken from a post url like https://t.me/c/CHANNEL_ID/1), the telegram account will join any new channels on setup", + "type": "json_loader", + }, + }, "description": """ The `TelethonExtractor` uses the Telethon library to archive posts and media from Telegram channels and groups. It supports private and public channels, downloading grouped posts with media, and can join channels using invite links @@ -46,5 +55,5 @@ To use the `TelethonExtractor`, you must configure the following: The first time you run, you will be prompted to do a authentication with the phone number associated, alternatively you can put your `anon.session` in the root. -""" -} \ No newline at end of file +""", +} diff --git a/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py b/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py index 65ea8cd..b06962e 100644 --- a/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py +++ b/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py @@ -1,12 +1,18 @@ - import shutil from telethon.sync import TelegramClient from telethon.errors import ChannelInvalidError from telethon.tl.functions.messages import ImportChatInviteRequest -from telethon.errors.rpcerrorlist import UserAlreadyParticipantError, FloodWaitError, InviteRequestSentError, InviteHashExpiredError +from telethon.errors.rpcerrorlist import ( + UserAlreadyParticipantError, + FloodWaitError, + InviteRequestSentError, + InviteHashExpiredError, +) from loguru import logger from tqdm import tqdm -import re, time, os +import re +import time +import os from auto_archiver.core import Extractor from auto_archiver.core import Metadata, Media @@ -17,9 +23,7 @@ class TelethonExtractor(Extractor): valid_url = re.compile(r"https:\/\/t\.me(\/c){0,1}\/(.+)\/(\d+)") invite_pattern = re.compile(r"t.me(\/joinchat){0,1}\/\+?(.+)") - def setup(self) -> None: - """ 1. makes a copy of session_file that is removed in cleanup 2. trigger login process for telegram or proceed if already saved in a session file @@ -34,7 +38,7 @@ class TelethonExtractor(Extractor): # initiate the client self.client = TelegramClient(self.session_file, self.api_id, self.api_hash) - + with self.client.start(): logger.success(f"SETUP {self.name} login works.") @@ -52,18 +56,20 @@ class TelethonExtractor(Extractor): channel_invite = self.channel_invites[i] channel_id = channel_invite.get("id", False) invite = channel_invite["invite"] - if (match := self.invite_pattern.search(invite)): + if match := self.invite_pattern.search(invite): try: if channel_id: ent = self.client.get_entity(int(channel_id)) # fails if not a member else: ent = self.client.get_entity(invite) # fails if not a member - logger.warning(f"please add the property id='{ent.id}' to the 'channel_invites' configuration where {invite=}, not doing so can lead to a minutes-long setup time due to telegram's rate limiting.") - except ValueError as e: + logger.warning( + f"please add the property id='{ent.id}' to the 'channel_invites' configuration where {invite=}, not doing so can lead to a minutes-long setup time due to telegram's rate limiting." + ) + except ValueError: logger.info(f"joining new channel {invite=}") try: self.client(ImportChatInviteRequest(match.group(2))) - except UserAlreadyParticipantError as e: + except UserAlreadyParticipantError: logger.info(f"already joined {invite=}") except InviteRequestSentError: logger.warning(f"already sent a join request with {invite} still no answer") @@ -95,7 +101,8 @@ class TelethonExtractor(Extractor): # detect URLs that we definitely cannot handle match = self.valid_url.search(url) logger.debug(f"TELETHON: {match=}") - if not match: return False + if not match: + return False is_private = match.group(1) == "/c" chat = int(match.group(2)) if is_private else match.group(2) @@ -105,45 +112,53 @@ class TelethonExtractor(Extractor): # NB: not using bot_token since then private channels cannot be archived: self.client.start(bot_token=self.bot_token) with self.client.start(): - # with self.client.start(bot_token=self.bot_token): + # with self.client.start(bot_token=self.bot_token): try: post = self.client.get_messages(chat, ids=post_id) except ValueError as e: logger.error(f"Could not fetch telegram {url} possibly it's private: {e}") return False except ChannelInvalidError as e: - logger.error(f"Could not fetch telegram {url}. This error may be fixed if you setup a bot_token in addition to api_id and api_hash (but then private channels will not be archived, we need to update this logic to handle both): {e}") + logger.error( + f"Could not fetch telegram {url}. This error may be fixed if you setup a bot_token in addition to api_id and api_hash (but then private channels will not be archived, we need to update this logic to handle both): {e}" + ) return False logger.debug(f"TELETHON GOT POST {post=}") - if post is None: return False + if post is None: + return False media_posts = self._get_media_posts_in_group(chat, post) - logger.debug(f'got {len(media_posts)=} for {url=}') + logger.debug(f"got {len(media_posts)=} for {url=}") tmp_dir = self.tmp_dir group_id = post.grouped_id if post.grouped_id is not None else post.id title = post.message for mp in media_posts: - if len(mp.message) > len(title): title = mp.message # save the longest text found (usually only 1) + if len(mp.message) > len(title): + title = mp.message # save the longest text found (usually only 1) # media can also be in entities if mp.entities: - other_media_urls = [e.url for e in mp.entities if hasattr(e, "url") and e.url and self._guess_file_type(e.url) in ["video", "image", "audio"]] + other_media_urls = [ + e.url + for e in mp.entities + if hasattr(e, "url") and e.url and self._guess_file_type(e.url) in ["video", "image", "audio"] + ] if len(other_media_urls): logger.debug(f"Got {len(other_media_urls)} other media urls from {mp.id=}: {other_media_urls}") for i, om_url in enumerate(other_media_urls): - filename = self.download_from_url(om_url, f'{chat}_{group_id}_{i}') + filename = self.download_from_url(om_url, f"{chat}_{group_id}_{i}") result.add_media(Media(filename=filename), id=f"{group_id}_{i}") - filename_dest = os.path.join(tmp_dir, f'{chat}_{group_id}', str(mp.id)) + filename_dest = os.path.join(tmp_dir, f"{chat}_{group_id}", str(mp.id)) filename = self.client.download_media(mp.media, filename_dest) if not filename: logger.debug(f"Empty media found, skipping {str(mp)=}") continue result.add_media(Media(filename)) - + result.set_title(title).set_timestamp(post.date).set("api_data", post.to_dict()) if post.message != title: result.set_content(post.message) diff --git a/src/auto_archiver/modules/thumbnail_enricher/__manifest__.py b/src/auto_archiver/modules/thumbnail_enricher/__manifest__.py index 1bd23b5..b11d17b 100644 --- a/src/auto_archiver/modules/thumbnail_enricher/__manifest__.py +++ b/src/auto_archiver/modules/thumbnail_enricher/__manifest__.py @@ -2,18 +2,19 @@ "name": "Thumbnail Enricher", "type": ["enricher"], "requires_setup": False, - "dependencies": { - "python": ["loguru", "ffmpeg"], - "bin": ["ffmpeg"] - }, + "dependencies": {"python": ["loguru", "ffmpeg"], "bin": ["ffmpeg"]}, "configs": { - "thumbnails_per_minute": {"default": 60, - "type": "int", - "help": "how many thumbnails to generate per minute of video, can be limited by max_thumbnails"}, - "max_thumbnails": {"default": 16, - "type": "int", - "help": "limit the number of thumbnails to generate per video, 0 means no limit"}, + "thumbnails_per_minute": { + "default": 60, + "type": "int", + "help": "how many thumbnails to generate per minute of video, can be limited by max_thumbnails", }, + "max_thumbnails": { + "default": 16, + "type": "int", + "help": "limit the number of thumbnails to generate per video, 0 means no limit", + }, + }, "description": """ Generates thumbnails for video files to provide visual previews. @@ -27,5 +28,5 @@ - Requires `ffmpeg` to be installed and accessible via the system's PATH. - Handles videos without pre-existing duration metadata by probing with `ffmpeg`. - Skips enrichment for non-video media files. - """ + """, } diff --git a/src/auto_archiver/modules/thumbnail_enricher/thumbnail_enricher.py b/src/auto_archiver/modules/thumbnail_enricher/thumbnail_enricher.py index 8178cd8..1543cec 100644 --- a/src/auto_archiver/modules/thumbnail_enricher/thumbnail_enricher.py +++ b/src/auto_archiver/modules/thumbnail_enricher/thumbnail_enricher.py @@ -6,7 +6,9 @@ visual snapshots of the video's keyframes, helping users preview content and identify important moments without watching the entire video. """ -import ffmpeg, os + +import ffmpeg +import os from loguru import logger from auto_archiver.core import Enricher @@ -18,7 +20,7 @@ class ThumbnailEnricher(Enricher): """ Generates thumbnails for all the media """ - + def enrich(self, to_enrich: Metadata) -> None: """ Uses or reads the video duration to generate thumbnails @@ -36,7 +38,9 @@ class ThumbnailEnricher(Enricher): if duration is None: try: probe = ffmpeg.probe(m.filename) - duration = float(next(stream for stream in probe['streams'] if stream['codec_type'] == 'video')['duration']) + duration = float( + next(stream for stream in probe["streams"] if stream["codec_type"] == "video")["duration"] + ) to_enrich.media[m_id].set("duration", duration) except Exception as e: logger.error(f"error getting duration of video {m.filename}: {e}") @@ -48,11 +52,13 @@ class ThumbnailEnricher(Enricher): thumbnails_media = [] for index, timestamp in enumerate(timestamps): output_path = os.path.join(folder, f"out{index}.jpg") - ffmpeg.input(m.filename, ss=timestamp).filter('scale', 512, -1).output(output_path, vframes=1, loglevel="quiet").run() + ffmpeg.input(m.filename, ss=timestamp).filter("scale", 512, -1).output( + output_path, vframes=1, loglevel="quiet" + ).run() try: - thumbnails_media.append(Media( - filename=output_path) + thumbnails_media.append( + Media(filename=output_path) .set("id", f"thumbnail_{index}") .set("timestamp", "%.3fs" % timestamp) ) diff --git a/src/auto_archiver/modules/tiktok_tikwm_extractor/__init__.py b/src/auto_archiver/modules/tiktok_tikwm_extractor/__init__.py deleted file mode 100644 index 25a20f5..0000000 --- a/src/auto_archiver/modules/tiktok_tikwm_extractor/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .tiktok_tikwm_extractor import TiktokTikwmExtractor \ No newline at end of file diff --git a/src/auto_archiver/modules/tiktok_tikwm_extractor/__manifest__.py b/src/auto_archiver/modules/tiktok_tikwm_extractor/__manifest__.py deleted file mode 100644 index 56d8e3e..0000000 --- a/src/auto_archiver/modules/tiktok_tikwm_extractor/__manifest__.py +++ /dev/null @@ -1,23 +0,0 @@ -{ - "name": "Tiktok Tikwm Extractor", - "type": ["extractor"], - "requires_setup": False, - "dependencies": { - "python": ["loguru", "requests"], - "bin": [] - }, - "description": """ - Uses an unofficial TikTok video download platform's API to download videos: https://tikwm.com/ - - This extractor complements the generic_extractor which can already get TikTok videos, but this one can extract special videos like those marked as sensitive. - - ### Features - - Downloads the video and, if possible, also the video cover. - - Stores extra metadata about the post like author information, and more as returned by tikwm.com. - - ### Notes - - If tikwm.com is down, this extractor will not work. - - If tikwm.com changes their API, this extractor may break. - - If no video is found, this extractor will consider the extraction failed. - """ -} diff --git a/src/auto_archiver/modules/tiktok_tikwm_extractor/tiktok_tikwm_extractor.py b/src/auto_archiver/modules/tiktok_tikwm_extractor/tiktok_tikwm_extractor.py deleted file mode 100644 index 8b07775..0000000 --- a/src/auto_archiver/modules/tiktok_tikwm_extractor/tiktok_tikwm_extractor.py +++ /dev/null @@ -1,75 +0,0 @@ -import re -import requests -from loguru import logger -from datetime import datetime, timezone -from yt_dlp.extractor.tiktok import TikTokIE - -from auto_archiver.core import Extractor -from auto_archiver.core import Metadata, Media - - -class TiktokTikwmExtractor(Extractor): - """ - Extractor for TikTok that uses an unofficial API and can capture content that requires a login, like sensitive content. - """ - TIKWM_ENDPOINT = "https://www.tikwm.com/api/?url={url}" - - def download(self, item: Metadata) -> Metadata: - url = item.get_url() - - if not re.match(TikTokIE._VALID_URL, url): - return False - - endpoint = TiktokTikwmExtractor.TIKWM_ENDPOINT.format(url=url) - - r = requests.get(endpoint) - if r.status_code != 200: - logger.error(f"unexpected status code '{r.status_code}' from tikwm.com for {url=}:") - return False - - try: - json_response = r.json() - except ValueError: - logger.error(f"failed to parse JSON response from tikwm.com for {url=}") - return False - - if not json_response.get('msg') == 'success' or not (api_data := json_response.get('data', {})): - logger.error(f"failed to get a valid response from tikwm.com for {url=}: {json_response}") - return False - - # tries to get the non-watermarked version first - video_url = api_data.pop("play", api_data.pop("wmplay", None)) - if not video_url: - logger.error(f"no valid video URL found in response from tikwm.com for {url=}") - return False - - # prepare result, start by downloading video - result = Metadata() - - # get the cover if possible - cover_url = api_data.pop("origin_cover", api_data.pop("cover", api_data.pop("ai_dynamic_cover", None))) - if cover_url and (cover_downloaded := self.download_from_url(cover_url)): - result.add_media(Media(cover_downloaded)) - - # get the video or fail - video_downloaded = self.download_from_url(video_url, f"vid_{api_data.get('id', '')}") - if not video_downloaded: - logger.error(f"failed to download video from {video_url}") - return False - video_media = Media(video_downloaded) - if duration := api_data.pop("duration", None): - video_media.set("duration", duration) - result.add_media(video_media) - - # add remaining metadata - result.set_title(api_data.pop("title", "")) - - if created_at := api_data.pop("create_time", None): - result.set_timestamp(datetime.fromtimestamp(created_at, tz=timezone.utc)) - - if (author := api_data.pop("author", None)): - result.set("author", author) - - result.set("api_data", api_data) - - return result.success("tikwm") diff --git a/src/auto_archiver/modules/timestamping_enricher/__manifest__.py b/src/auto_archiver/modules/timestamping_enricher/__manifest__.py index 6ad9c57..c451437 100644 --- a/src/auto_archiver/modules/timestamping_enricher/__manifest__.py +++ b/src/auto_archiver/modules/timestamping_enricher/__manifest__.py @@ -3,38 +3,29 @@ "type": ["enricher"], "requires_setup": True, "dependencies": { - "python": [ - "loguru", - "slugify", - "tsp_client", - "asn1crypto", - "certvalidator", - "certifi" - ], + "python": ["loguru", "slugify", "tsp_client", "asn1crypto", "certvalidator", "certifi"], }, "configs": { "tsa_urls": { "default": [ - # [Adobe Approved Trust List] and [Windows Cert Store] - "http://timestamp.digicert.com", - "http://timestamp.identrust.com", - # "https://timestamp.entrust.net/TSS/RFC3161sha2TS", # not valid for timestamping - # "https://timestamp.sectigo.com", # wait 15 seconds between each request. - - # [Adobe: European Union Trusted Lists]. - # "https://timestamp.sectigo.com/qualified", # wait 15 seconds between each request. - - # [Windows Cert Store] - "http://timestamp.globalsign.com/tsa/r6advanced1", - # [Adobe: European Union Trusted Lists] and [Windows Cert Store] - # "http://ts.quovadisglobal.com/eu", # not valid for timestamping - # "http://tsa.belgium.be/connect", # self-signed certificate in certificate chain - # "https://timestamp.aped.gov.gr/qtss", # self-signed certificate in certificate chain - # "http://tsa.sep.bg", # self-signed certificate in certificate chain - # "http://tsa.izenpe.com", #unable to get local issuer certificate - # "http://kstamp.keynectis.com/KSign", # unable to get local issuer certificate - "http://tss.accv.es:8318/tsa", - ], + # [Adobe Approved Trust List] and [Windows Cert Store] + "http://timestamp.digicert.com", + "http://timestamp.identrust.com", + # "https://timestamp.entrust.net/TSS/RFC3161sha2TS", # not valid for timestamping + # "https://timestamp.sectigo.com", # wait 15 seconds between each request. + # [Adobe: European Union Trusted Lists]. + # "https://timestamp.sectigo.com/qualified", # wait 15 seconds between each request. + # [Windows Cert Store] + "http://timestamp.globalsign.com/tsa/r6advanced1", + # [Adobe: European Union Trusted Lists] and [Windows Cert Store] + # "http://ts.quovadisglobal.com/eu", # not valid for timestamping + # "http://tsa.belgium.be/connect", # self-signed certificate in certificate chain + # "https://timestamp.aped.gov.gr/qtss", # self-signed certificate in certificate chain + # "http://tsa.sep.bg", # self-signed certificate in certificate chain + # "http://tsa.izenpe.com", #unable to get local issuer certificate + # "http://kstamp.keynectis.com/KSign", # unable to get local issuer certificate + "http://tss.accv.es:8318/tsa", + ], "help": "List of RFC3161 Time Stamp Authorities to use, separate with commas if passed via the command line.", } }, @@ -50,5 +41,5 @@ ### Notes - Should be run after the `hash_enricher` to ensure file hashes are available. - Requires internet access to interact with the configured TSAs. - """ + """, } diff --git a/src/auto_archiver/modules/timestamping_enricher/timestamping_enricher.py b/src/auto_archiver/modules/timestamping_enricher/timestamping_enricher.py index b83e86c..586b7f8 100644 --- a/src/auto_archiver/modules/timestamping_enricher/timestamping_enricher.py +++ b/src/auto_archiver/modules/timestamping_enricher/timestamping_enricher.py @@ -11,6 +11,7 @@ import certifi from auto_archiver.core import Enricher from auto_archiver.core import Metadata, Media + class TimestampingEnricher(Enricher): """ Uses several RFC3161 Time Stamp Authorities to generate a timestamp token that will be preserved. This can be used to prove that a certain file existed at a certain time, useful for legal purposes, for example, to prove that a certain file was not tampered with after a certain date. @@ -25,7 +26,9 @@ class TimestampingEnricher(Enricher): logger.debug(f"RFC3161 timestamping existing files for {url=}") # create a new text file with the existing media hashes - hashes = [m.get("hash").replace("SHA-256:", "").replace("SHA3-512:", "") for m in to_enrich.media if m.get("hash")] + hashes = [ + m.get("hash").replace("SHA-256:", "").replace("SHA3-512:", "") for m in to_enrich.media if m.get("hash") + ] if not len(hashes): logger.warning(f"No hashes found in {url=}") @@ -35,17 +38,18 @@ class TimestampingEnricher(Enricher): hashes_fn = os.path.join(tmp_dir, "hashes.txt") data_to_sign = "\n".join(hashes) - with open(hashes_fn, "w") as f: + with open(hashes_fn, "w") as f: f.write(data_to_sign) hashes_media = Media(filename=hashes_fn) timestamp_tokens = [] from slugify import slugify + for tsa_url in self.tsa_urls: try: signing_settings = SigningSettings(tsp_server=tsa_url, digest_algorithm=DigestAlgorithm.SHA256) signer = TSPSigner() - message = bytes(data_to_sign, encoding='utf8') + message = bytes(data_to_sign, encoding="utf8") # send TSQ and get TSR from the TSA server signed = signer.sign(message=message, signing_settings=signing_settings) # fail if there's any issue with the certificates, uses certifi list of trusted CAs @@ -54,7 +58,8 @@ class TimestampingEnricher(Enricher): cert_chain = self.download_and_verify_certificate(signed) # continue with saving the timestamp token tst_fn = os.path.join(tmp_dir, f"timestamp_token_{slugify(tsa_url)}") - with open(tst_fn, "wb") as f: f.write(signed) + with open(tst_fn, "wb") as f: + f.write(signed) timestamp_tokens.append(Media(filename=tst_fn).set("tsa", tsa_url).set("cert_chain", cert_chain)) except Exception as e: logger.warning(f"Error while timestamping {url=} with {tsa_url=}: {e}") @@ -75,7 +80,7 @@ class TimestampingEnricher(Enricher): tst = ContentInfo.load(signed) trust_roots = [] - with open(certifi.where(), 'rb') as f: + with open(certifi.where(), "rb") as f: for _, _, der_bytes in pem.unarmor(f.read(), multiple=True): trust_roots.append(der_bytes) context = ValidationContext(trust_roots=trust_roots) @@ -83,11 +88,11 @@ class TimestampingEnricher(Enricher): certificates = tst["content"]["certificates"] first_cert = certificates[0].dump() intermediate_certs = [] - for i in range(1, len(certificates)): # cannot use list comprehension [1:] + for i in range(1, len(certificates)): # cannot use list comprehension [1:] intermediate_certs.append(certificates[i].dump()) validator = CertificateValidator(first_cert, intermediate_certs=intermediate_certs, validation_context=context) - path = validator.validate_usage({'digital_signature'}, extended_key_usage={'time_stamping'}) + path = validator.validate_usage({"digital_signature"}, extended_key_usage={"time_stamping"}) cert_chain = [] for cert in path: @@ -96,4 +101,4 @@ class TimestampingEnricher(Enricher): f.write(cert.dump()) cert_chain.append(Media(filename=cert_fn).set("subject", cert.subject.native["common_name"])) - return cert_chain \ No newline at end of file + return cert_chain diff --git a/src/auto_archiver/modules/twitter_api_extractor/__init__.py b/src/auto_archiver/modules/twitter_api_extractor/__init__.py index 7005965..54e7b6c 100644 --- a/src/auto_archiver/modules/twitter_api_extractor/__init__.py +++ b/src/auto_archiver/modules/twitter_api_extractor/__init__.py @@ -1 +1 @@ -from .twitter_api_extractor import TwitterApiExtractor \ No newline at end of file +from .twitter_api_extractor import TwitterApiExtractor diff --git a/src/auto_archiver/modules/twitter_api_extractor/__manifest__.py b/src/auto_archiver/modules/twitter_api_extractor/__manifest__.py index 05d1ac0..203155f 100644 --- a/src/auto_archiver/modules/twitter_api_extractor/__manifest__.py +++ b/src/auto_archiver/modules/twitter_api_extractor/__manifest__.py @@ -3,21 +3,28 @@ "type": ["extractor"], "requires_setup": True, "dependencies": { - "python": ["requests", - "loguru", - "pytwitter", - "slugify",], - "bin": [""] + "python": [ + "requests", + "loguru", + "pytwitter", + "slugify", + ], + "bin": [""], }, "configs": { - "bearer_token": {"default": None, "help": "[deprecated: see bearer_tokens] twitter API bearer_token which is enough for archiving, if not provided you will need consumer_key, consumer_secret, access_token, access_secret"}, - "bearer_tokens": {"default": [], "help": " a list of twitter API bearer_token which is enough for archiving, if not provided you will need consumer_key, consumer_secret, access_token, access_secret, if provided you can still add those for better rate limits. CSV of bearer tokens if provided via the command line", - }, - "consumer_key": {"default": None, "help": "twitter API consumer_key"}, - "consumer_secret": {"default": None, "help": "twitter API consumer_secret"}, - "access_token": {"default": None, "help": "twitter API access_token"}, - "access_secret": {"default": None, "help": "twitter API access_secret"}, + "bearer_token": { + "default": None, + "help": "[deprecated: see bearer_tokens] twitter API bearer_token which is enough for archiving, if not provided you will need consumer_key, consumer_secret, access_token, access_secret", }, + "bearer_tokens": { + "default": [], + "help": " a list of twitter API bearer_token which is enough for archiving, if not provided you will need consumer_key, consumer_secret, access_token, access_secret, if provided you can still add those for better rate limits. CSV of bearer tokens if provided via the command line", + }, + "consumer_key": {"default": None, "help": "twitter API consumer_key"}, + "consumer_secret": {"default": None, "help": "twitter API consumer_secret"}, + "access_token": {"default": None, "help": "twitter API access_token"}, + "access_secret": {"default": None, "help": "twitter API access_secret"}, + }, "description": """ The `TwitterApiExtractor` fetches tweets and associated media using the Twitter API. It supports multiple API configurations for extended rate limits and reliable access. @@ -39,6 +46,5 @@ - **Access Token and Secret**: Complements the consumer key for enhanced API capabilities. Credentials can be obtained by creating a Twitter developer account at [Twitter Developer Platform](https://developer.twitter.com/en). - """ -, + """, } diff --git a/src/auto_archiver/modules/twitter_api_extractor/twitter_api_extractor.py b/src/auto_archiver/modules/twitter_api_extractor/twitter_api_extractor.py index 72fd2f2..1c08235 100644 --- a/src/auto_archiver/modules/twitter_api_extractor/twitter_api_extractor.py +++ b/src/auto_archiver/modules/twitter_api_extractor/twitter_api_extractor.py @@ -11,8 +11,8 @@ from slugify import slugify from auto_archiver.core import Extractor from auto_archiver.core import Metadata, Media -class TwitterApiExtractor(Extractor): +class TwitterApiExtractor(Extractor): valid_url: re.Pattern = re.compile(r"(?:twitter|x).com\/(?:\#!\/)?(\w+)\/status(?:es)?\/(\d+)") def setup(self) -> None: @@ -23,30 +23,38 @@ class TwitterApiExtractor(Extractor): if self.bearer_token: self.apis.append(Api(bearer_token=self.bearer_token)) if self.consumer_key and self.consumer_secret and self.access_token and self.access_secret: - self.apis.append(Api(consumer_key=self.consumer_key, consumer_secret=self.consumer_secret, - access_token=self.access_token, access_secret=self.access_secret)) - assert self.api_client is not None, "Missing Twitter API configurations, please provide either AND/OR (consumer_key, consumer_secret, access_token, access_secret) to use this archiver, you can provide both for better rate-limit results." + self.apis.append( + Api( + consumer_key=self.consumer_key, + consumer_secret=self.consumer_secret, + access_token=self.access_token, + access_secret=self.access_secret, + ) + ) + assert self.api_client is not None, ( + "Missing Twitter API configurations, please provide either AND/OR (consumer_key, consumer_secret, access_token, access_secret) to use this archiver, you can provide both for better rate-limit results." + ) @property # getter .mimetype def api_client(self) -> str: return self.apis[self.api_index] - + def sanitize_url(self, url: str) -> str: # expand URL if t.co and clean tracker GET params - if 'https://t.co/' in url: + if "https://t.co/" in url: try: r = requests.get(url, timeout=30) - logger.debug(f'Expanded url {url} to {r.url}') + logger.debug(f"Expanded url {url} to {r.url}") url = r.url - except: - logger.error(f'Failed to expand url {url}') + except Exception: + logger.error(f"Failed to expand url {url}") return url - def download(self, item: Metadata) -> Metadata: # call download retry until success or no more apis while self.api_index < len(self.apis): - if res := self.download_retry(item): return res + if res := self.download_retry(item): + return res self.api_index += 1 self.api_index = 0 return False @@ -54,7 +62,8 @@ class TwitterApiExtractor(Extractor): def get_username_tweet_id(self, url): # detect URLs that we definitely cannot handle matches = self.valid_url.findall(url) - if not len(matches): return False, False + if not len(matches): + return False, False username, tweet_id = matches[0] # only one URL supported logger.debug(f"Found {username=} and {tweet_id=} in {url=}") @@ -65,10 +74,16 @@ class TwitterApiExtractor(Extractor): url = item.get_url() # detect URLs that we definitely cannot handle username, tweet_id = self.get_username_tweet_id(url) - if not username: return False + if not username: + return False try: - tweet = self.api_client.get_tweet(tweet_id, expansions=["attachments.media_keys"], media_fields=["type", "duration_ms", "url", "variants"], tweet_fields=["attachments", "author_id", "created_at", "entities", "id", "text", "possibly_sensitive"]) + tweet = self.api_client.get_tweet( + tweet_id, + expansions=["attachments.media_keys"], + media_fields=["type", "duration_ms", "url", "variants"], + tweet_fields=["attachments", "author_id", "created_at", "entities", "id", "text", "possibly_sensitive"], + ) logger.debug(tweet) except Exception as e: logger.error(f"Could not get tweet: {e}") @@ -88,29 +103,35 @@ class TwitterApiExtractor(Extractor): mimetype = "image/jpeg" elif hasattr(m, "variants"): variant = self.choose_variant(m.variants) - if not variant: continue + if not variant: + continue media.set("src", variant.url) mimetype = variant.content_type else: continue logger.info(f"Found media {media}") ext = mimetypes.guess_extension(mimetype) - media.filename = self.download_from_url(media.get("src"), f'{slugify(url)}_{i}{ext}') + media.filename = self.download_from_url(media.get("src"), f"{slugify(url)}_{i}{ext}") result.add_media(media) - result.set_content(json.dumps({ - "id": tweet.data.id, - "text": tweet.data.text, - "created_at": tweet.data.created_at, - "author_id": tweet.data.author_id, - "geo": tweet.data.geo, - "lang": tweet.data.lang, - "media": urls - }, ensure_ascii=False, indent=4)) + result.set_content( + json.dumps( + { + "id": tweet.data.id, + "text": tweet.data.text, + "created_at": tweet.data.created_at, + "author_id": tweet.data.author_id, + "geo": tweet.data.geo, + "lang": tweet.data.lang, + "media": urls, + }, + ensure_ascii=False, + indent=4, + ) + ) return result.success("twitter-api") def choose_variant(self, variants): - """ Chooses the highest quality variable possible out of a list of variants """ diff --git a/src/auto_archiver/modules/vk_extractor/__manifest__.py b/src/auto_archiver/modules/vk_extractor/__manifest__.py index 61e454e..ed16331 100644 --- a/src/auto_archiver/modules/vk_extractor/__manifest__.py +++ b/src/auto_archiver/modules/vk_extractor/__manifest__.py @@ -7,10 +7,8 @@ "python": ["loguru", "vk_url_scraper"], }, "configs": { - "username": {"required": True, - "help": "valid VKontakte username"}, - "password": {"required": True, - "help": "valid VKontakte password"}, + "username": {"required": True, "help": "valid VKontakte username"}, + "password": {"required": True, "help": "valid VKontakte password"}, "session_file": { "default": "secrets/vk_config.v2.json", "help": "valid VKontakte password", diff --git a/src/auto_archiver/modules/vk_extractor/vk_extractor.py b/src/auto_archiver/modules/vk_extractor/vk_extractor.py index 99527c4..997b0a8 100644 --- a/src/auto_archiver/modules/vk_extractor/vk_extractor.py +++ b/src/auto_archiver/modules/vk_extractor/vk_extractor.py @@ -7,7 +7,7 @@ from auto_archiver.core import Metadata, Media class VkExtractor(Extractor): - """" + """ " VK videos are handled by YTDownloader, this archiver gets posts text and images. Currently only works for /wall posts """ @@ -18,11 +18,13 @@ class VkExtractor(Extractor): def download(self, item: Metadata) -> Metadata: url = item.get_url() - if "vk.com" not in item.netloc: return False + if "vk.com" not in item.netloc: + return False # some urls can contain multiple wall/photo/... parts and all will be fetched vk_scrapes = self.vks.scrape(url) - if not len(vk_scrapes): return False + if not len(vk_scrapes): + return False logger.debug(f"VK: got {len(vk_scrapes)} scraped instances") result = Metadata() diff --git a/src/auto_archiver/modules/wacz_extractor_enricher/__manifest__.py b/src/auto_archiver/modules/wacz_extractor_enricher/__manifest__.py index 9b373b9..97e3bf6 100644 --- a/src/auto_archiver/modules/wacz_extractor_enricher/__manifest__.py +++ b/src/auto_archiver/modules/wacz_extractor_enricher/__manifest__.py @@ -4,34 +4,38 @@ "entry_point": "wacz_extractor_enricher::WaczExtractorEnricher", "requires_setup": True, "dependencies": { - "python": [ - "loguru", - "jsonlines", - "warcio" - ], + "python": ["loguru", "jsonlines", "warcio"], # TODO? - "bin": [ - "docker" - ] + "bin": ["docker"], }, "configs": { - "profile": {"default": None, "help": "browsertrix-profile (for profile generation see https://github.com/webrecorder/browsertrix-crawler#creating-and-using-browser-profiles)."}, - "docker_commands": {"default": None, "help":"if a custom docker invocation is needed"}, - "timeout": {"default": 120, - "type": "int", - "help": "timeout for WACZ generation in seconds", "type": "int"}, - "extract_media": {"default": False, - "type": 'bool', - "help": "If enabled all the images/videos/audio present in the WACZ archive will be extracted into separate Media and appear in the html report. The .wacz file will be kept untouched." - }, - "extract_screenshot": {"default": True, - "type": 'bool', - "help": "If enabled the screenshot captured by browsertrix will be extracted into separate Media and appear in the html report. The .wacz file will be kept untouched." - }, - "socks_proxy_host": {"default": None, "help": "SOCKS proxy host for browsertrix-crawler, use in combination with socks_proxy_port. eg: user:password@host"}, - "socks_proxy_port": {"default": None, "type":"int", "help": "SOCKS proxy port for browsertrix-crawler, use in combination with socks_proxy_host. eg 1234"}, - "proxy_server": {"default": None, "help": "SOCKS server proxy URL, in development"}, + "profile": { + "default": None, + "help": "browsertrix-profile (for profile generation see https://github.com/webrecorder/browsertrix-crawler#creating-and-using-browser-profiles).", }, + "docker_commands": {"default": None, "help": "if a custom docker invocation is needed"}, + "timeout": {"default": 120, "help": "timeout for WACZ generation in seconds", "type": "int"}, + "extract_media": { + "default": False, + "type": "bool", + "help": "If enabled all the images/videos/audio present in the WACZ archive will be extracted into separate Media and appear in the html report. The .wacz file will be kept untouched.", + }, + "extract_screenshot": { + "default": True, + "type": "bool", + "help": "If enabled the screenshot captured by browsertrix will be extracted into separate Media and appear in the html report. The .wacz file will be kept untouched.", + }, + "socks_proxy_host": { + "default": None, + "help": "SOCKS proxy host for browsertrix-crawler, use in combination with socks_proxy_port. eg: user:password@host", + }, + "socks_proxy_port": { + "default": None, + "type": "int", + "help": "SOCKS proxy port for browsertrix-crawler, use in combination with socks_proxy_host. eg 1234", + }, + "proxy_server": {"default": None, "help": "SOCKS server proxy URL, in development"}, + }, "description": """ Creates .WACZ archives of web pages using the `browsertrix-crawler` tool, with options for media extraction and screenshot saving. [Browsertrix-crawler](https://crawler.docs.browsertrix.com/user-guide/) is a headless browser-based crawler that archives web pages in WACZ format. @@ -45,5 +49,5 @@ ### Notes - Requires Docker for running `browsertrix-crawler` . - Configurable via parameters for timeout, media extraction, screenshots, and proxy settings. - """ + """, } diff --git a/src/auto_archiver/modules/wacz_extractor_enricher/wacz_extractor_enricher.py b/src/auto_archiver/modules/wacz_extractor_enricher/wacz_extractor_enricher.py index ff7314a..975d49a 100644 --- a/src/auto_archiver/modules/wacz_extractor_enricher/wacz_extractor_enricher.py +++ b/src/auto_archiver/modules/wacz_extractor_enricher/wacz_extractor_enricher.py @@ -1,6 +1,8 @@ import jsonlines import mimetypes -import os, shutil, subprocess +import os +import shutil +import subprocess from zipfile import ZipFile from loguru import logger from warcio.archiveiterator import ArchiveIterator @@ -19,13 +21,12 @@ class WaczExtractorEnricher(Enricher, Extractor): """ def setup(self) -> None: - - self.use_docker = os.environ.get('WACZ_ENABLE_DOCKER') or not os.environ.get('RUNNING_IN_DOCKER') - self.docker_in_docker = os.environ.get('WACZ_ENABLE_DOCKER') and os.environ.get('RUNNING_IN_DOCKER') + self.use_docker = os.environ.get("WACZ_ENABLE_DOCKER") or not os.environ.get("RUNNING_IN_DOCKER") + self.docker_in_docker = os.environ.get("WACZ_ENABLE_DOCKER") and os.environ.get("RUNNING_IN_DOCKER") self.cwd_dind = f"/crawls/crawls{random_str(8)}" - self.browsertrix_home_host = os.environ.get('BROWSERTRIX_HOME_HOST') - self.browsertrix_home_container = os.environ.get('BROWSERTRIX_HOME_CONTAINER') or self.browsertrix_home_host + self.browsertrix_home_host = os.environ.get("BROWSERTRIX_HOME_HOST") + self.browsertrix_home_container = os.environ.get("BROWSERTRIX_HOME_CONTAINER") or self.browsertrix_home_host # create crawls folder if not exists, so it can be safely removed in cleanup if self.docker_in_docker: os.makedirs(self.cwd_dind, exist_ok=True) @@ -55,21 +56,32 @@ class WaczExtractorEnricher(Enricher, Extractor): cmd = [ "crawl", - "--url", url, - "--scopeType", "page", + "--url", + url, + "--scopeType", + "page", "--generateWACZ", - "--text", "to-pages", - "--screenshot", "fullPage", - "--collection", collection, - "--id", collection, - "--saveState", "never", - "--behaviors", "autoscroll,autoplay,autofetch,siteSpecific", - "--behaviorTimeout", str(self.timeout), - "--timeout", str(self.timeout), - "--diskUtilization", "99", + "--text", + "to-pages", + "--screenshot", + "fullPage", + "--collection", + collection, + "--id", + collection, + "--saveState", + "never", + "--behaviors", + "autoscroll,autoplay,autofetch,siteSpecific", + "--behaviorTimeout", + str(self.timeout), + "--timeout", + str(self.timeout), + "--diskUtilization", + "99", # "--blockAds" # note: this has been known to cause issues on cloudflare protected sites ] - + if self.docker_in_docker: cmd.extend(["--cwd", self.cwd_dind]) @@ -80,7 +92,14 @@ class WaczExtractorEnricher(Enricher, Extractor): if self.docker_commands: cmd = self.docker_commands + cmd else: - cmd = ["docker", "run", "--rm", "-v", f"{browsertrix_home_host}:/crawls/", "webrecorder/browsertrix-crawler"] + cmd + cmd = [ + "docker", + "run", + "--rm", + "-v", + f"{browsertrix_home_host}:/crawls/", + "webrecorder/browsertrix-crawler", + ] + cmd if self.profile: profile_fn = os.path.join(browsertrix_home_container, "profile.tar.gz") @@ -109,7 +128,6 @@ class WaczExtractorEnricher(Enricher, Extractor): logger.error(f"WACZ generation failed: {e}") return False - if self.docker_in_docker: wacz_fn = os.path.join(self.cwd_dind, "collections", collection, f"{collection}.wacz") elif self.use_docker: @@ -138,11 +156,10 @@ class WaczExtractorEnricher(Enricher, Extractor): logger.info(f"Parsing pages.jsonl {jsonl_fn=}") with jsonlines.open(jsonl_fn) as reader: for obj in reader: - if 'title' in obj: - to_enrich.set_title(obj['title']) - if 'text' in obj: - to_enrich.set_content(obj['text']) - + if "title" in obj: + to_enrich.set_title(obj["title"]) + if "text" in obj: + to_enrich.set_content(obj["text"]) return True @@ -155,36 +172,41 @@ class WaczExtractorEnricher(Enricher, Extractor): # unzipping the .wacz tmp_dir = self.tmp_dir unzipped_dir = os.path.join(tmp_dir, "unzipped") - with ZipFile(wacz_filename, 'r') as z_obj: + with ZipFile(wacz_filename, "r") as z_obj: z_obj.extractall(path=unzipped_dir) # if warc is split into multiple gzip chunks, merge those warc_dir = os.path.join(unzipped_dir, "archive") warc_filename = os.path.join(tmp_dir, "merged.warc") - with open(warc_filename, 'wb') as outfile: + with open(warc_filename, "wb") as outfile: for filename in sorted(os.listdir(warc_dir)): - if filename.endswith('.gz'): + if filename.endswith(".gz"): chunk_file = os.path.join(warc_dir, filename) - with open(chunk_file, 'rb') as infile: + with open(chunk_file, "rb") as infile: shutil.copyfileobj(infile, outfile) # get media out of .warc counter = 0 seen_urls = set() - import json - with open(warc_filename, 'rb') as warc_stream: + + with open(warc_filename, "rb") as warc_stream: for record in ArchiveIterator(warc_stream): # only include fetched resources - if record.rec_type == "resource" and record.content_type == "image/png" and self.extract_screenshot: # screenshots + if ( + record.rec_type == "resource" and record.content_type == "image/png" and self.extract_screenshot + ): # screenshots fn = os.path.join(tmp_dir, f"warc-file-{counter}.png") - with open(fn, "wb") as outf: outf.write(record.raw_stream.read()) + with open(fn, "wb") as outf: + outf.write(record.raw_stream.read()) m = Media(filename=fn) to_enrich.add_media(m, "browsertrix-screenshot") counter += 1 - if not self.extract_media: continue + if not self.extract_media: + continue - if record.rec_type != 'response': continue - record_url = record.rec_headers.get_header('WARC-Target-URI') + if record.rec_type != "response": + continue + record_url = record.rec_headers.get_header("WARC-Target-URI") if not UrlUtil.is_relevant_url(record_url): logger.debug(f"Skipping irrelevant URL {record_url} but it's still present in the WACZ.") continue @@ -194,8 +216,10 @@ class WaczExtractorEnricher(Enricher, Extractor): # filter by media mimetypes content_type = record.http_headers.get("Content-Type") - if not content_type: continue - if not any(x in content_type for x in ["video", "image", "audio"]): continue + if not content_type: + continue + if not any(x in content_type for x in ["video", "image", "audio"]): + continue # create local file and add media ext = mimetypes.guess_extension(content_type) @@ -203,7 +227,8 @@ class WaczExtractorEnricher(Enricher, Extractor): fn = os.path.join(tmp_dir, warc_fn) record_url_best_qual = UrlUtil.twitter_best_quality_url(record_url) - with open(fn, "wb") as outf: outf.write(record.raw_stream.read()) + with open(fn, "wb") as outf: + outf.write(record.raw_stream.read()) m = Media(filename=fn) m.set("src", record_url) @@ -213,12 +238,16 @@ class WaczExtractorEnricher(Enricher, Extractor): m.filename = self.download_from_url(record_url_best_qual, warc_fn) m.set("src", record_url_best_qual) m.set("src_alternative", record_url) - except Exception as e: logger.warning(f"Unable to download best quality URL for {record_url=} got error {e}, using original in WARC.") + except Exception as e: + logger.warning( + f"Unable to download best quality URL for {record_url=} got error {e}, using original in WARC." + ) # remove bad videos - if m.is_video() and not m.is_valid_video(): continue - + if m.is_video() and not m.is_valid_video(): + continue + to_enrich.add_media(m, warc_fn) counter += 1 seen_urls.add(record_url) - logger.info(f"WACZ extract_media/extract_screenshot finished, found {counter} relevant media file(s)") \ No newline at end of file + logger.info(f"WACZ extract_media/extract_screenshot finished, found {counter} relevant media file(s)") diff --git a/src/auto_archiver/modules/wayback_extractor_enricher/__init__.py b/src/auto_archiver/modules/wayback_extractor_enricher/__init__.py index b69332d..b6fb182 100644 --- a/src/auto_archiver/modules/wayback_extractor_enricher/__init__.py +++ b/src/auto_archiver/modules/wayback_extractor_enricher/__init__.py @@ -1 +1 @@ -from .wayback_extractor_enricher import WaybackExtractorEnricher \ No newline at end of file +from .wayback_extractor_enricher import WaybackExtractorEnricher diff --git a/src/auto_archiver/modules/wayback_extractor_enricher/wayback_extractor_enricher.py b/src/auto_archiver/modules/wayback_extractor_enricher/wayback_extractor_enricher.py index 1763b12..f06effd 100644 --- a/src/auto_archiver/modules/wayback_extractor_enricher/wayback_extractor_enricher.py +++ b/src/auto_archiver/modules/wayback_extractor_enricher/wayback_extractor_enricher.py @@ -1,16 +1,18 @@ import json from loguru import logger -import time, requests +import time +import requests from auto_archiver.core import Extractor, Enricher from auto_archiver.utils import url as UrlUtil from auto_archiver.core import Metadata + class WaybackExtractorEnricher(Enricher, Extractor): """ Submits the current URL to the webarchive and returns a job_id or completed archive. - The Wayback machine will rate-limit IP heavy usage. + The Wayback machine will rate-limit IP heavy usage. """ def download(self, item: Metadata) -> Metadata: @@ -22,8 +24,10 @@ class WaybackExtractorEnricher(Enricher, Extractor): def enrich(self, to_enrich: Metadata) -> bool: proxies = {} - if self.proxy_http: proxies["http"] = self.proxy_http - if self.proxy_https: proxies["https"] = self.proxy_https + if self.proxy_http: + proxies["http"] = self.proxy_http + if self.proxy_https: + proxies["https"] = self.proxy_https url = to_enrich.get_url() if UrlUtil.is_auth_wall(url): @@ -36,15 +40,12 @@ class WaybackExtractorEnricher(Enricher, Extractor): logger.info(f"Wayback enricher had already been executed: {to_enrich.get('wayback')}") return True - ia_headers = { - "Accept": "application/json", - "Authorization": f"LOW {self.key}:{self.secret}" - } - post_data = {'url': url} + ia_headers = {"Accept": "application/json", "Authorization": f"LOW {self.key}:{self.secret}"} + post_data = {"url": url} if self.if_not_archived_within: post_data["if_not_archived_within"] = self.if_not_archived_within # see https://docs.google.com/document/d/1Nsv52MvSjbLb2PCpHlat0gkzw0EvtSgpKHu4mk0MnrA for more options - r = requests.post('https://web.archive.org/save/', headers=ia_headers, data=post_data, proxies=proxies) + r = requests.post("https://web.archive.org/save/", headers=ia_headers, data=post_data, proxies=proxies) if r.status_code != 200: logger.error(em := f"Internet archive failed with status of {r.status_code}: {r.json()}") @@ -53,15 +54,14 @@ class WaybackExtractorEnricher(Enricher, Extractor): # check job status try: - job_id = r.json().get('job_id') + job_id = r.json().get("job_id") if not job_id: logger.error(f"Wayback failed with {r.json()}") return False - except json.decoder.JSONDecodeError as e: + except json.decoder.JSONDecodeError: logger.error(f"Expected a JSON with job_id from Wayback and got {r.text}") return False - # waits at most timeout seconds until job is completed, otherwise only enriches the job_id information start_time = time.time() wayback_url = False @@ -69,17 +69,19 @@ class WaybackExtractorEnricher(Enricher, Extractor): while not wayback_url and time.time() - start_time <= self.timeout: try: logger.debug(f"GETting status for {job_id=} on {url=} ({attempt=})") - r_status = requests.get(f'https://web.archive.org/save/status/{job_id}', headers=ia_headers, proxies=proxies) + r_status = requests.get( + f"https://web.archive.org/save/status/{job_id}", headers=ia_headers, proxies=proxies + ) r_json = r_status.json() - if r_status.status_code == 200 and r_json['status'] == 'success': + if r_status.status_code == 200 and r_json["status"] == "success": wayback_url = f"https://web.archive.org/web/{r_json['timestamp']}/{r_json['original_url']}" - elif r_status.status_code != 200 or r_json['status'] != 'pending': + elif r_status.status_code != 200 or r_json["status"] != "pending": logger.error(f"Wayback failed with {r_json}") return False except requests.exceptions.RequestException as e: logger.warning(f"RequestException: fetching status for {url=} due to: {e}") break - except json.decoder.JSONDecodeError as e: + except json.decoder.JSONDecodeError: logger.error(f"Expected a JSON from Wayback and got {r.text} for {url=}") break except Exception as e: @@ -91,6 +93,8 @@ class WaybackExtractorEnricher(Enricher, Extractor): if wayback_url: to_enrich.set("wayback", wayback_url) else: - to_enrich.set("wayback", {"job_id": job_id, "check_status": f'https://web.archive.org/save/status/{job_id}'}) + to_enrich.set( + "wayback", {"job_id": job_id, "check_status": f"https://web.archive.org/save/status/{job_id}"} + ) to_enrich.set("check wayback", f"https://web.archive.org/web/*/{url}") return True diff --git a/src/auto_archiver/modules/whisper_enricher/__init__.py b/src/auto_archiver/modules/whisper_enricher/__init__.py index d3d3526..d69bdd1 100644 --- a/src/auto_archiver/modules/whisper_enricher/__init__.py +++ b/src/auto_archiver/modules/whisper_enricher/__init__.py @@ -1 +1 @@ -from .whisper_enricher import WhisperEnricher \ No newline at end of file +from .whisper_enricher import WhisperEnricher diff --git a/src/auto_archiver/modules/whisper_enricher/__manifest__.py b/src/auto_archiver/modules/whisper_enricher/__manifest__.py index 0e09d03..e7af7d1 100644 --- a/src/auto_archiver/modules/whisper_enricher/__manifest__.py +++ b/src/auto_archiver/modules/whisper_enricher/__manifest__.py @@ -6,19 +6,26 @@ "python": ["s3_storage", "loguru", "requests"], }, "configs": { - "api_endpoint": {"required": True, - "help": "WhisperApi api endpoint, eg: https://whisperbox-api.com/api/v1, a deployment of https://github.com/bellingcat/whisperbox-transcribe."}, - "api_key": {"required": True, - "help": "WhisperApi api key for authentication"}, - "include_srt": {"default": False, - "type": "bool", - "help": "Whether to include a subtitle SRT (SubRip Subtitle file) for the video (can be used in video players)."}, - "timeout": {"default": 90, - "type": "int", - "help": "How many seconds to wait at most for a successful job completion."}, - "action": {"default": "translate", - "help": "which Whisper operation to execute", - "choices": ["transcribe", "translate", "language_detection"]}, + "api_endpoint": { + "required": True, + "help": "WhisperApi api endpoint, eg: https://whisperbox-api.com/api/v1, a deployment of https://github.com/bellingcat/whisperbox-transcribe.", + }, + "api_key": {"required": True, "help": "WhisperApi api key for authentication"}, + "include_srt": { + "default": False, + "type": "bool", + "help": "Whether to include a subtitle SRT (SubRip Subtitle file) for the video (can be used in video players).", + }, + "timeout": { + "default": 90, + "type": "int", + "help": "How many seconds to wait at most for a successful job completion.", + }, + "action": { + "default": "translate", + "help": "which Whisper operation to execute", + "choices": ["transcribe", "translate", "language_detection"], + }, }, "description": """ Integrates with a Whisper API service to transcribe, translate, or detect the language of audio and video files. @@ -35,5 +42,5 @@ - Only compatible with S3-compatible storage systems for media file accessibility. - ** This stores the media files in S3 prior to enriching them as Whisper requires public URLs to access the media files. - Handles multiple jobs and retries for failed or incomplete processing. - """ + """, } diff --git a/src/auto_archiver/modules/whisper_enricher/whisper_enricher.py b/src/auto_archiver/modules/whisper_enricher/whisper_enricher.py index d63d2ed..063bd26 100644 --- a/src/auto_archiver/modules/whisper_enricher/whisper_enricher.py +++ b/src/auto_archiver/modules/whisper_enricher/whisper_enricher.py @@ -1,10 +1,12 @@ import traceback -import requests, time +import requests +import time from loguru import logger from auto_archiver.core import Enricher from auto_archiver.core import Metadata, Media + class WhisperEnricher(Enricher): """ Connects with a Whisper API service to get texts out of audio @@ -13,15 +15,15 @@ class WhisperEnricher(Enricher): """ def setup(self) -> None: - self.stores = self.config['steps']['storages'] + self.stores = self.config["steps"]["storages"] self.s3 = self.module_factory.get_module("s3_storage", self.config) - if not "s3_storage" in self.stores: - logger.error("WhisperEnricher: To use the WhisperEnricher you need to use S3Storage so files are accessible publicly to the whisper service being called.") + if "s3_storage" not in self.stores: + logger.error( + "WhisperEnricher: To use the WhisperEnricher you need to use S3Storage so files are accessible publicly to the whisper service being called." + ) return - def enrich(self, to_enrich: Metadata) -> None: - url = to_enrich.get_url() logger.debug(f"WHISPER[{self.action}]: iterating media items for {url=}.") @@ -36,28 +38,33 @@ class WhisperEnricher(Enricher): logger.debug(f"JOB SUBMITTED: {job_id=} for {m.key=}") to_enrich.media[i].set("whisper_model", {"job_id": job_id}) except Exception as e: - logger.error(f"Failed to submit whisper job for {m.filename=} with error {e}\n{traceback.format_exc()}") + logger.error( + f"Failed to submit whisper job for {m.filename=} with error {e}\n{traceback.format_exc()}" + ) job_results = self.check_jobs(job_results) for i, m in enumerate(to_enrich.media): if m.is_video() or m.is_audio(): job_id = to_enrich.media[i].get("whisper_model", {}).get("job_id") - if not job_id: continue - to_enrich.media[i].set("whisper_model", { - "job_id": job_id, - "job_status_check": f"{self.api_endpoint}/jobs/{job_id}", - "job_artifacts_check": f"{self.api_endpoint}/jobs/{job_id}/artifacts", - **(job_results[job_id] if job_results[job_id] else {"result": "incomplete or failed job"}) - }) + if not job_id: + continue + to_enrich.media[i].set( + "whisper_model", + { + "job_id": job_id, + "job_status_check": f"{self.api_endpoint}/jobs/{job_id}", + "job_artifacts_check": f"{self.api_endpoint}/jobs/{job_id}/artifacts", + **(job_results[job_id] if job_results[job_id] else {"result": "incomplete or failed job"}), + }, + ) # append the extracted text to the content of the post so it gets written to the DBs like gsheets text column if job_results[job_id]: - for k,v in job_results[job_id].items(): + for k, v in job_results[job_id].items(): if "_text" in k and len(v): to_enrich.set_content(f"\n[automatic video transcript]: {v}") def submit_job(self, media: Media): - s3_url = self.s3.get_cdn_url(media) assert s3_url in media.urls, f"Could not find S3 url ({s3_url}) in list of stored media urls " payload = { @@ -66,10 +73,14 @@ class WhisperEnricher(Enricher): # "language": "string" # may be a config } logger.debug(f"calling API with {payload=}") - response = requests.post(f'{self.api_endpoint}/jobs', json=payload, headers={'Authorization': f'Bearer {self.api_key}'}) - assert response.status_code == 201, f"calling the whisper api {self.api_endpoint} returned a non-success code: {response.status_code}" + response = requests.post( + f"{self.api_endpoint}/jobs", json=payload, headers={"Authorization": f"Bearer {self.api_key}"} + ) + assert response.status_code == 201, ( + f"calling the whisper api {self.api_endpoint} returned a non-success code: {response.status_code}" + ) logger.debug(response.json()) - return response.json()['id'] + return response.json()["id"] def check_jobs(self, job_results: dict): start_time = time.time() @@ -77,37 +88,50 @@ class WhisperEnricher(Enricher): while not all_completed and (time.time() - start_time) <= self.timeout: all_completed = True for job_id in job_results: - if job_results[job_id] != False: continue + if job_results[job_id] is not False: + continue all_completed = False # at least one not ready - try: job_results[job_id] = self.check_job(job_id) + try: + job_results[job_id] = self.check_job(job_id) except Exception as e: logger.error(f"Failed to check {job_id=} with error {e}\n{traceback.format_exc()}") - if not all_completed: time.sleep(3) + if not all_completed: + time.sleep(3) return job_results def check_job(self, job_id): - r = requests.get(f'{self.api_endpoint}/jobs/{job_id}', headers={'Authorization': f'Bearer {self.api_key}'}) + r = requests.get(f"{self.api_endpoint}/jobs/{job_id}", headers={"Authorization": f"Bearer {self.api_key}"}) assert r.status_code == 200, f"Job status did not respond with 200, instead with: {r.status_code}" j = r.json() logger.debug(f"Checked job {job_id=} with status='{j['status']}'") - if j['status'] == "processing": return False - elif j['status'] == "error": return f"Error: {j['meta']['error']}" - elif j['status'] == "success": - r_res = requests.get(f'{self.api_endpoint}/jobs/{job_id}/artifacts', headers={'Authorization': f'Bearer {self.api_key}'}) - assert r_res.status_code == 200, f"Job artifacts did not respond with 200, instead with: {r_res.status_code}" + if j["status"] == "processing": + return False + elif j["status"] == "error": + return f"Error: {j['meta']['error']}" + elif j["status"] == "success": + r_res = requests.get( + f"{self.api_endpoint}/jobs/{job_id}/artifacts", headers={"Authorization": f"Bearer {self.api_key}"} + ) + assert r_res.status_code == 200, ( + f"Job artifacts did not respond with 200, instead with: {r_res.status_code}" + ) logger.success(r_res.json()) result = {} for art_id, artifact in enumerate(r_res.json()): subtitle = [] full_text = [] for i, d in enumerate(artifact.get("data")): - subtitle.append(f"{i+1}\n{d.get('start')} --> {d.get('end')}\n{d.get('text').strip()}") - full_text.append(d.get('text').strip()) - if not len(subtitle): continue - if self.include_srt: result[f"artifact_{art_id}_subtitle"] = "\n".join(subtitle) + subtitle.append(f"{i + 1}\n{d.get('start')} --> {d.get('end')}\n{d.get('text').strip()}") + full_text.append(d.get("text").strip()) + if not len(subtitle): + continue + if self.include_srt: + result[f"artifact_{art_id}_subtitle"] = "\n".join(subtitle) result[f"artifact_{art_id}_text"] = "\n".join(full_text) # call /delete endpoint on timely success - r_del = requests.delete(f'{self.api_endpoint}/jobs/{job_id}', headers={'Authorization': f'Bearer {self.api_key}'}) + r_del = requests.delete( + f"{self.api_endpoint}/jobs/{job_id}", headers={"Authorization": f"Bearer {self.api_key}"} + ) logger.debug(f"DELETE whisper {job_id=} result: {r_del.status_code}") return result return False diff --git a/src/auto_archiver/utils/__init__.py b/src/auto_archiver/utils/__init__.py index 46ca191..a8fa77f 100644 --- a/src/auto_archiver/utils/__init__.py +++ b/src/auto_archiver/utils/__init__.py @@ -1,7 +1,8 @@ -""" Auto Archiver Utilities. """ +"""Auto Archiver Utilities.""" + # we need to explicitly expose the available imports here from .misc import * from .webdriver import Webdriver # handy utils from ytdlp -from yt_dlp.utils import (clean_html, traverse_obj, strip_or_none, url_or_none) \ No newline at end of file +from yt_dlp.utils import clean_html, traverse_obj, strip_or_none, url_or_none diff --git a/src/auto_archiver/utils/misc.py b/src/auto_archiver/utils/misc.py index 49ef0b5..fe1864b 100644 --- a/src/auto_archiver/utils/misc.py +++ b/src/auto_archiver/utils/misc.py @@ -16,22 +16,23 @@ def mkdir_if_not_exists(folder): def expand_url(url): # expand short URL links - if 'https://t.co/' in url: + if "https://t.co/" in url: try: r = requests.get(url) - logger.debug(f'Expanded url {url} to {r.url}') + logger.debug(f"Expanded url {url} to {r.url}") return r.url - except: - logger.error(f'Failed to expand url {url}') + except Exception: + logger.error(f"Failed to expand url {url}") return url def getattr_or(o: object, prop: str, default=None): try: res = getattr(o, prop) - if res is None: raise + if res is None: + raise return res - except: + except Exception: return default @@ -61,18 +62,19 @@ def random_str(length: int = 32) -> str: return str(uuid.uuid4()).replace("-", "")[:length] -def calculate_file_hash(filename: str, hash_algo = hashlib.sha256, chunksize: int = 16000000) -> str: +def calculate_file_hash(filename: str, hash_algo=hashlib.sha256, chunksize: int = 16000000) -> str: hash = hash_algo() with open(filename, "rb") as f: while True: buf = f.read(chunksize) - if not buf: break + if not buf: + break hash.update(buf) return hash.hexdigest() def get_datetime_from_str(dt_str: str, fmt: str | None = None, dayfirst=True) -> datetime | None: - """ parse a datetime string with option of passing a specific format + """parse a datetime string with option of passing a specific format Args: dt_str: the datetime string to parse @@ -88,19 +90,24 @@ def get_datetime_from_str(dt_str: str, fmt: str | None = None, dayfirst=True) -> def get_timestamp(ts, utc=True, iso=True, dayfirst=True) -> str | datetime | None: - """ Consistent parsing of timestamps. + """Consistent parsing of timestamps. Args: If utc=True, the timezone is set to UTC, if iso=True, the output is an iso string Use dayfirst to signify between date formats which put the date vs month first: e.g. DD/MM/YYYY vs MM/DD/YYYY - """ - if not ts: return + """ + if not ts: + return try: - if isinstance(ts, str): ts = parse_dt(ts, dayfirst=dayfirst) - if isinstance(ts, (int, float)): ts = datetime.fromtimestamp(ts) - if utc: ts = ts.replace(tzinfo=timezone.utc) - if iso: return ts.isoformat() + if isinstance(ts, str): + ts = parse_dt(ts, dayfirst=dayfirst) + if isinstance(ts, (int, float)): + ts = datetime.fromtimestamp(ts) + if utc: + ts = ts.replace(tzinfo=timezone.utc) + if iso: + return ts.isoformat() return ts except Exception as e: logger.error(f"Unable to parse timestamp {ts}: {e}") diff --git a/src/auto_archiver/utils/url.py b/src/auto_archiver/utils/url.py index 061f4aa..169ed87 100644 --- a/src/auto_archiver/utils/url.py +++ b/src/auto_archiver/utils/url.py @@ -4,8 +4,8 @@ from ipaddress import ip_address AUTHWALL_URLS = [ - re.compile(r"https:\/\/t\.me(\/c)\/(.+)\/(\d+)"), # telegram private channels - re.compile(r"https:\/\/www\.instagram\.com"), # instagram + re.compile(r"https:\/\/t\.me(\/c)\/(.+)\/(\d+)"), # telegram private channels + re.compile(r"https:\/\/www\.instagram\.com"), # instagram ] @@ -14,17 +14,16 @@ def check_url_or_raise(url: str) -> bool | ValueError: Blocks localhost, private, reserved, and link-local IPs and all non-http/https schemes. """ - if not (url.startswith("http://") or url.startswith("https://")): raise ValueError(f"Invalid URL scheme for url {url}") - + parsed = urlparse(url) if not parsed.hostname: raise ValueError(f"Invalid URL hostname for url {url}") - + if parsed.hostname == "localhost": raise ValueError(f"Localhost URLs cannot be parsed for security reasons (for url {url})") - + if parsed.scheme not in ["http", "https"]: raise ValueError(f"Invalid URL scheme, only http and https supported (for url {url})") @@ -32,7 +31,7 @@ def check_url_or_raise(url: str) -> bool | ValueError: ip = ip_address(parsed.hostname) except ValueError: pass - + else: if not ip.is_global: raise ValueError(f"IP address {ip} is not globally reachable") @@ -42,18 +41,21 @@ def check_url_or_raise(url: str) -> bool | ValueError: raise ValueError(f"Link-local IP address {ip} used") if ip.is_private: raise ValueError(f"Private IP address {ip} used") - + return True + def domain_for_url(url: str) -> str: """ SECURITY: parse the domain using urllib to avoid any potential security issues """ return urlparse(url).netloc + def clean(url: str) -> str: return url + def is_auth_wall(url: str) -> bool: """ checks if URL is behind an authentication wall meaning steps like wayback, wacz, ... may not work @@ -64,13 +66,15 @@ def is_auth_wall(url: str) -> bool: return False + def remove_get_parameters(url: str) -> str: # http://example.com/file.mp4?t=1 -> http://example.com/file.mp4 # useful for mimetypes to work parsed_url = urlparse(url) - new_url = urlunparse(parsed_url._replace(query='')) + new_url = urlunparse(parsed_url._replace(query="")) return new_url + def is_relevant_url(url: str) -> bool: """ Detect if a detected media URL is recurring and therefore irrelevant to a specific archive. Useful, for example, for the enumeration of the media files in WARC files which include profile pictures, favicons, etc. @@ -78,42 +82,59 @@ def is_relevant_url(url: str) -> bool: clean_url = remove_get_parameters(url) # favicons - if "favicon" in url: return False + if "favicon" in url: + return False # ifnore icons - if clean_url.endswith(".ico"): return False + if clean_url.endswith(".ico"): + return False # ignore SVGs - if remove_get_parameters(url).endswith(".svg"): return False + if remove_get_parameters(url).endswith(".svg"): + return False # twitter profile pictures - if "twimg.com/profile_images" in url: return False - if "twimg.com" in url and "/default_profile_images" in url: return False + if "twimg.com/profile_images" in url: + return False + if "twimg.com" in url and "/default_profile_images" in url: + return False # instagram profile pictures - if "https://scontent.cdninstagram.com/" in url and "150x150" in url: return False + if "https://scontent.cdninstagram.com/" in url and "150x150" in url: + return False # instagram recurring images - if "https://static.cdninstagram.com/rsrc.php/" in url: return False + if "https://static.cdninstagram.com/rsrc.php/" in url: + return False # telegram - if "https://telegram.org/img/emoji/" in url: return False + if "https://telegram.org/img/emoji/" in url: + return False # youtube - if "https://www.youtube.com/s/gaming/emoji/" in url: return False - if "https://yt3.ggpht.com" in url and "default-user=" in url: return False - if "https://www.youtube.com/s/search/audio/" in url: return False + if "https://www.youtube.com/s/gaming/emoji/" in url: + return False + if "https://yt3.ggpht.com" in url and "default-user=" in url: + return False + if "https://www.youtube.com/s/search/audio/" in url: + return False # ok - if " https://ok.ru/res/i/" in url: return False + if " https://ok.ru/res/i/" in url: + return False # vk - if "https://vk.com/emoji/" in url: return False - if "vk.com/images/" in url: return False - if "vk.com/images/reaction/" in url: return False + if "https://vk.com/emoji/" in url: + return False + if "vk.com/images/" in url: + return False + if "vk.com/images/reaction/" in url: + return False # wikipedia - if "wikipedia.org/static" in url: return False + if "wikipedia.org/static" in url: + return False return True + def twitter_best_quality_url(url: str) -> str: """ some twitter image URLs point to a less-than best quality diff --git a/src/auto_archiver/utils/webdriver.py b/src/auto_archiver/utils/webdriver.py index cb4e2a9..2ba185e 100644 --- a/src/auto_archiver/utils/webdriver.py +++ b/src/auto_archiver/utils/webdriver.py @@ -1,10 +1,11 @@ -""" This Webdriver class acts as a context manager for the selenium webdriver. """ +"""This Webdriver class acts as a context manager for the selenium webdriver.""" + from __future__ import annotations import os import time -#import domain_for_url +# import domain_for_url from urllib.parse import urlparse, urlunparse from http.cookiejar import MozillaCookieJar @@ -19,16 +20,15 @@ from loguru import logger class CookieSettingDriver(webdriver.Firefox): - facebook_accept_cookies: bool cookies: str cookiejar: MozillaCookieJar def __init__(self, cookies, cookiejar, facebook_accept_cookies, *args, **kwargs): - if os.environ.get('RUNNING_IN_DOCKER'): + if os.environ.get("RUNNING_IN_DOCKER"): # Selenium doesn't support linux-aarch64 driver, we need to set this manually - kwargs['service'] = webdriver.FirefoxService(executable_path='/usr/local/bin/geckodriver') - + kwargs["service"] = webdriver.FirefoxService(executable_path="/usr/local/bin/geckodriver") + super(CookieSettingDriver, self).__init__(*args, **kwargs) self.cookies = cookies self.cookiejar = cookiejar @@ -38,42 +38,43 @@ class CookieSettingDriver(webdriver.Firefox): if self.cookies or self.cookiejar: # set up the driver to make it not 'cookie averse' (needs a context/URL) # get the 'robots.txt' file which should be quick and easy - robots_url = urlunparse(urlparse(url)._replace(path='/robots.txt', query='', fragment='')) + robots_url = urlunparse(urlparse(url)._replace(path="/robots.txt", query="", fragment="")) super(CookieSettingDriver, self).get(robots_url) if self.cookies: # an explicit cookie is set for this site, use that first for cookie in self.cookies.split(";"): for name, value in cookie.split("="): - self.driver.add_cookie({'name': name, 'value': value}) + self.driver.add_cookie({"name": name, "value": value}) elif self.cookiejar: domain = urlparse(url).netloc.lstrip("www.") for cookie in self.cookiejar: if domain in cookie.domain: try: - self.add_cookie({ - 'name': cookie.name, - 'value': cookie.value, - 'path': cookie.path, - 'domain': cookie.domain, - 'secure': bool(cookie.secure), - 'expiry': cookie.expires - }) + self.add_cookie( + { + "name": cookie.name, + "value": cookie.value, + "path": cookie.path, + "domain": cookie.domain, + "secure": bool(cookie.secure), + "expiry": cookie.expires, + } + ) except Exception as e: logger.warning(f"Failed to add cookie to webdriver: {e}") - + if self.facebook_accept_cookies: try: - logger.debug(f'Trying fb click accept cookie popup.') + logger.debug("Trying fb click accept cookie popup.") super(CookieSettingDriver, self).get("http://www.facebook.com") essential_only = self.find_element(By.XPATH, "//span[contains(text(), 'Decline optional cookies')]") essential_only.click() - logger.debug(f'fb click worked') + logger.debug("fb click worked") # linux server needs a sleep otherwise facebook cookie won't have worked and we'll get a popup on next page time.sleep(2) except Exception as e: - logger.warning(f'Failed on fb accept cookies.', e) - + logger.warning("Failed on fb accept cookies.", e) # now get the actual URL super(CookieSettingDriver, self).get(url) @@ -87,9 +88,14 @@ class CookieSettingDriver(webdriver.Firefox): pass else: - # for all other sites, try and use some common button text to reject/accept cookies - for text in ["Refuse non-essential cookies", "Decline optional cookies", "Reject additional cookies", "Reject all", "Accept all cookies"]: + for text in [ + "Refuse non-essential cookies", + "Decline optional cookies", + "Reject additional cookies", + "Reject all", + "Accept all cookies", + ]: try: xpath = f"//*[contains(translate(text(), 'ABCDEFGHIJKLMNOPQRSTUVWXYZ', 'abcdefghijklmnopqrstuvwxyz'), '{text.lower()}')]" WebDriverWait(self, 5).until(EC.element_to_be_clickable((By.XPATH, xpath))).click() @@ -97,11 +103,18 @@ class CookieSettingDriver(webdriver.Firefox): except selenium_exceptions.WebDriverException: pass - + class Webdriver: - def __init__(self, width: int, height: int, timeout_seconds: int, - facebook_accept_cookies: bool = False, http_proxy: str = "", - print_options: dict = {}, auth: dict = {}) -> webdriver: + def __init__( + self, + width: int, + height: int, + timeout_seconds: int, + facebook_accept_cookies: bool = False, + http_proxy: str = "", + print_options: dict = {}, + auth: dict = {}, + ) -> webdriver: self.width = width self.height = height self.timeout_seconds = timeout_seconds @@ -116,23 +129,29 @@ class Webdriver: def __enter__(self) -> webdriver: options = webdriver.FirefoxOptions() options.add_argument("--headless") - options.add_argument(f'--proxy-server={self.http_proxy}') - options.set_preference('network.protocol-handler.external.tg', False) + options.add_argument(f"--proxy-server={self.http_proxy}") + options.set_preference("network.protocol-handler.external.tg", False) # if facebook cookie popup is present, force the browser to English since then it's easier to click the 'Decline optional cookies' option if self.facebook_accept_cookies: - options.add_argument('--lang=en') + options.add_argument("--lang=en") try: - self.driver = CookieSettingDriver(cookies=self.auth.get('cookies'), cookiejar=self.auth.get('cookies_jar'), - facebook_accept_cookies=self.facebook_accept_cookies, options=options) + self.driver = CookieSettingDriver( + cookies=self.auth.get("cookies"), + cookiejar=self.auth.get("cookies_jar"), + facebook_accept_cookies=self.facebook_accept_cookies, + options=options, + ) self.driver.set_window_size(self.width, self.height) self.driver.set_page_load_timeout(self.timeout_seconds) self.driver.print_options = self.print_options except selenium_exceptions.TimeoutException as e: - logger.error(f"failed to get new webdriver, possibly due to insufficient system resources or timeout settings: {e}") + logger.error( + f"failed to get new webdriver, possibly due to insufficient system resources or timeout settings: {e}" + ) return self.driver - + def __exit__(self, exc_type, exc_val, exc_tb): self.driver.close() self.driver.quit() diff --git a/src/auto_archiver/version.py b/src/auto_archiver/version.py index dd700b6..d0e6a5f 100644 --- a/src/auto_archiver/version.py +++ b/src/auto_archiver/version.py @@ -1,7 +1,8 @@ -""" Version information for the auto_archiver package. - TODO: This is a placeholder to replicate previous versioning. +"""Version information for the auto_archiver package. +TODO: This is a placeholder to replicate previous versioning. """ + from importlib.metadata import version as get_version VERSION_SHORT = get_version("auto_archiver") @@ -9,4 +10,4 @@ VERSION_SHORT = get_version("auto_archiver") # This is mainly for nightly builds which have the suffix ".dev$DATE". See # https://semver.org/#is-v123-a-semantic-version for the semantics. _SUFFIX = "" -__version__ = f"{VERSION_SHORT}{_SUFFIX}" \ No newline at end of file +__version__ = f"{VERSION_SHORT}{_SUFFIX}" diff --git a/tests/conftest.py b/tests/conftest.py index a94abcd..379bfc2 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,6 +1,7 @@ """ pytest conftest file, for shared fixtures and configuration """ + import os import pickle from datetime import datetime, timezone @@ -16,32 +17,36 @@ from auto_archiver.core.module import ModuleFactory # that you only want to run if everything else succeeds (e.g. API calls). The order here is important # what comes first will be run first (at the end of all other tests not mentioned) # format is the name of the module (python file) without the .py extension -TESTS_TO_RUN_LAST = ['test_twitter_api_archiver'] +TESTS_TO_RUN_LAST = ["test_twitter_api_archiver"] + @pytest.fixture def setup_module(request): - def _setup_module(module_name, config={}): - + def _setup_module(module_name, config=None): + if config is None: + config = {} module_factory = ModuleFactory() if isinstance(module_name, type): # get the module name: # if the class does not have a .name, use the name of the parent folder - module_name = module_name.__module__.rsplit(".",2)[-2] + module_name = module_name.__module__.rsplit(".", 2)[-2] m = module_factory.get_module(module_name, {module_name: config}) # add the tmp_dir to the module tmp_dir = TemporaryDirectory() m.tmp_dir = tmp_dir.name - + def cleanup(): tmp_dir.cleanup() + request.addfinalizer(cleanup) return m return _setup_module + @pytest.fixture def check_hash(): def _check_hash(filename: str, hash: str): @@ -51,6 +56,7 @@ def check_hash(): return _check_hash + @pytest.fixture def make_item(): def _make_item(url: str, **kwargs) -> Metadata: @@ -62,7 +68,6 @@ def make_item(): return _make_item - def pytest_collection_modifyitems(items): module_mapping = {item: item.module.__name__.split(".")[-1] for item in items} @@ -78,13 +83,13 @@ def pytest_collection_modifyitems(items): items[:] = sorted_items - # Incremental testing - fail tests in a class if any previous test fails # taken from https://docs.pytest.org/en/latest/example/simple.html#incremental-testing-test-steps # store history of failures per test class name and per index in parametrize (if parametrize used) _test_failed_incremental: Dict[str, Dict[Tuple[int, ...], str]] = {} + def pytest_runtest_makereport(item, call): if "incremental" in item.keywords: # incremental marker is used @@ -93,17 +98,11 @@ def pytest_runtest_makereport(item, call): # retrieve the class name of the test cls_name = str(item.cls) # retrieve the index of the test (if parametrize is used in combination with incremental) - parametrize_index = ( - tuple(item.callspec.indices.values()) - if hasattr(item, "callspec") - else () - ) + parametrize_index = tuple(item.callspec.indices.values()) if hasattr(item, "callspec") else () # retrieve the name of the test function test_name = item.originalname or item.name # store in _test_failed_incremental the original name of the failed test - _test_failed_incremental.setdefault(cls_name, {}).setdefault( - parametrize_index, test_name - ) + _test_failed_incremental.setdefault(cls_name, {}).setdefault(parametrize_index, test_name) def pytest_runtest_setup(item): @@ -119,16 +118,17 @@ def pytest_runtest_setup(item): pytest.xfail(f"previous test failed ({test_name})") - @pytest.fixture() def unpickle(): """ Returns a helper function that unpickles a file ** gets the file from the test_files directory: tests/data/ ** """ + def _unpickle(path): with open(os.path.join("tests/data", path), "rb") as f: return pickle.load(f) + return _unpickle @@ -145,9 +145,9 @@ def sample_datetime(): return datetime(2023, 1, 1, 12, 0, tzinfo=timezone.utc) -@pytest.fixture(autouse=True) +@pytest.fixture def mock_sleep(mocker): - """Globally mock time.sleep to avoid delays.""" + """Mock time.sleep to avoid delays.""" return mocker.patch("time.sleep") @@ -156,4 +156,4 @@ def metadata(): metadata = Metadata() metadata.set("_processed_at", "2021-01-01T00:00:00") metadata.set_url("https://example.com") - return metadata \ No newline at end of file + return metadata diff --git a/tests/data/dropin.py b/tests/data/dropin.py index 0049c48..93c7500 100644 --- a/tests/data/dropin.py +++ b/tests/data/dropin.py @@ -1,5 +1,6 @@ # this is a dummy class used to test importing a dropin in the # generic extractor by filename/path + class Dropin: - pass \ No newline at end of file + pass diff --git a/tests/data/test_modules/example_module/__init__.py b/tests/data/test_modules/example_module/__init__.py index 560a9b9..02986fc 100644 --- a/tests/data/test_modules/example_module/__init__.py +++ b/tests/data/test_modules/example_module/__init__.py @@ -1 +1 @@ -from .example_module import ExampleModule \ No newline at end of file +from .example_module import ExampleModule diff --git a/tests/data/test_modules/example_module/__manifest__.py b/tests/data/test_modules/example_module/__manifest__.py index e3a26bb..064b3c9 100644 --- a/tests/data/test_modules/example_module/__manifest__.py +++ b/tests/data/test_modules/example_module/__manifest__.py @@ -16,14 +16,14 @@ "dependencies": { "python": ["loguru"], "bin": ["bash"], - }, - # configurations that this module takes. These are argparse-compliant dicationaries, that are + }, + # configurations that this module takes. These are argparse-compliant dicationaries, that are # used to create command line arguments when the programme is run. # The full name of the config option will become: `module_name.config_name` "configs": { - "csv_file": {"default": "db.csv", "help": "CSV file name"}, - "required_field": {"required": True, "help": "required field in the CSV file"}, - }, + "csv_file": {"default": "db.csv", "help": "CSV file name"}, + "required_field": {"required": True, "help": "required field in the CSV file"}, + }, # A description of the module, used for documentation "description": "This is an example module", -} \ No newline at end of file +} diff --git a/tests/data/test_modules/example_module/example_module.py b/tests/data/test_modules/example_module/example_module.py index 7def054..392abe0 100644 --- a/tests/data/test_modules/example_module/example_module.py +++ b/tests/data/test_modules/example_module/example_module.py @@ -1,5 +1,6 @@ from auto_archiver.core import Extractor, Enricher, Feeder, Database, Storage, Formatter, Metadata + class ExampleModule(Extractor, Enricher, Feeder, Database, Storage, Formatter): def download(self, item): print("download") @@ -7,7 +8,6 @@ class ExampleModule(Extractor, Enricher, Feeder, Database, Storage, Formatter): def __iter__(self): yield Metadata().set_url("https://example.com") - def done(self, result): print("done") @@ -16,13 +16,12 @@ class ExampleModule(Extractor, Enricher, Feeder, Database, Storage, Formatter): def get_cdn_url(self, media): return "nice_url" - + def save(self, item): print("save") - + def uploadf(self, file, key, **kwargs): print("uploadf") - def format(self, item): print("format") diff --git a/tests/databases/test_api_db.py b/tests/databases/test_api_db.py index 5d1ea84..2e87a87 100644 --- a/tests/databases/test_api_db.py +++ b/tests/databases/test_api_db.py @@ -1,6 +1,5 @@ import pytest -from auto_archiver.core import Metadata from auto_archiver.modules.api_db import AAApiDb @@ -41,9 +40,16 @@ def test_fetch(api_db, metadata, mocker): mock_datetime = mocker.patch("auto_archiver.core.metadata.datetime.datetime") mock_datetime.now.return_value = "2021-01-01T00:00:00" mock_get.return_value.status_code = 200 - mock_get.return_value.json.return_value = [{"result": {}}, {"result": - {'media': [], 'metadata': {'_processed_at': '2021-01-01T00:00:00', 'url': 'https://example.com'}, - 'status': 'no archiver'}}] + mock_get.return_value.json.return_value = [ + {"result": {}}, + { + "result": { + "media": [], + "metadata": {"_processed_at": "2021-01-01T00:00:00", "url": "https://example.com"}, + "status": "no archiver", + } + }, + ] assert api_db.fetch(metadata) == metadata @@ -52,8 +58,15 @@ def test_done_success(api_db, metadata, mocker): mock_post.return_value.status_code = 201 api_db.done(metadata) mock_post.assert_called_once() - mock_post.assert_called_once_with("https://api.example.com/interop/submit-archive", - json={'author_id': 'Someone', 'url': 'https://example.com', - 'public': False, 'group_id': '123', 'tags': ['[', ']'], 'result': '{"status": "no archiver", "metadata": {"_processed_at": "2021-01-01T00:00:00", "url": "https://example.com"}, "media": []}'}, - headers={'Authorization': 'Bearer test-token'}) - + mock_post.assert_called_once_with( + "https://api.example.com/interop/submit-archive", + json={ + "author_id": "Someone", + "url": "https://example.com", + "public": False, + "group_id": "123", + "tags": ["[", "]"], + "result": '{"status": "no archiver", "metadata": {"_processed_at": "2021-01-01T00:00:00", "url": "https://example.com"}, "media": []}', + }, + headers={"Authorization": "Bearer test-token"}, + ) diff --git a/tests/databases/test_atlos_db.py b/tests/databases/test_atlos_db.py index a73f1df..15f4b55 100644 --- a/tests/databases/test_atlos_db.py +++ b/tests/databases/test_atlos_db.py @@ -49,10 +49,8 @@ def test_failed_with_atlos_id(atlos_db, metadata, mocker): fake_resp = FakeAPIResponse({}, raise_error=False) post_mock = mocker.patch.object(atlos_db, "_post", return_value=fake_resp) atlos_db.failed(metadata, "failure reason") - expected_endpoint = f"/api/v2/source_material/metadata/42/auto_archiver" - expected_json = { - "metadata": {"processed": True, "status": "error", "error": "failure reason"} - } + expected_endpoint = "/api/v2/source_material/metadata/42/auto_archiver" + expected_json = {"metadata": {"processed": True, "status": "error", "error": "failure reason"}} post_mock.assert_called_once_with(expected_endpoint, json=expected_json) @@ -87,7 +85,7 @@ def test_done_with_atlos_id(atlos_db, metadata, mocker): fake_resp = FakeAPIResponse({}, raise_error=False) post_mock = mocker.patch.object(atlos_db, "_post", return_value=fake_resp) atlos_db.done(metadata) - expected_endpoint = f"/api/v2/source_material/metadata/99/auto_archiver" + expected_endpoint = "/api/v2/source_material/metadata/99/auto_archiver" expected_results = metadata.metadata.copy() expected_results["timestamp"] = now.isoformat() expected_json = { diff --git a/tests/databases/test_csv_db.py b/tests/databases/test_csv_db.py index afca0d8..bf5b7cb 100644 --- a/tests/databases/test_csv_db.py +++ b/tests/databases/test_csv_db.py @@ -1,4 +1,3 @@ - from auto_archiver.modules.csv_db import CSVDb from auto_archiver.core import Metadata @@ -9,12 +8,21 @@ def test_store_item(tmp_path, setup_module): temp_db = tmp_path / "temp_db.csv" db = setup_module(CSVDb, {"csv_file": temp_db.as_posix()}) - item = Metadata().set_url("http://example.com").set_title("Example").set_content("Example content").success("my-archiver") + item = ( + Metadata() + .set_url("http://example.com") + .set_title("Example") + .set_content("Example content") + .success("my-archiver") + ) db.done(item) with open(temp_db, "r", encoding="utf-8") as f: - assert f.read().strip() == f"status,metadata,media\nmy-archiver: success,\"{{'_processed_at': {repr(item.get('_processed_at'))}, 'url': 'http://example.com', 'title': 'Example', 'content': 'Example content'}}\",[]" + assert ( + f.read().strip() + == f"status,metadata,media\nmy-archiver: success,\"{{'_processed_at': {repr(item.get('_processed_at'))}, 'url': 'http://example.com', 'title': 'Example', 'content': 'Example content'}}\",[]" + ) # TODO: csv db doesn't have a fetch method - need to add it (?) - # assert db.fetch(item) == item \ No newline at end of file + # assert db.fetch(item) == item diff --git a/tests/databases/test_gsheet_db.py b/tests/databases/test_gsheet_db.py index 2f1202d..0760c79 100644 --- a/tests/databases/test_gsheet_db.py +++ b/tests/databases/test_gsheet_db.py @@ -28,6 +28,7 @@ def mock_metadata(mocker): metadata.get_first_image.return_value = None return metadata + @pytest.fixture def metadata(): metadata = Metadata() @@ -51,15 +52,31 @@ def mock_media(mocker): mock_media.get.return_value = "not-calculated" return mock_media + @pytest.fixture -def gsheets_db(mock_gworksheet, setup_module, mocker): +def gsheets_db(mock_gworksheet, setup_module, mocker) -> GsheetsFeederDB: mocker.patch("gspread.service_account") config: dict = { "sheet": "testsheet", "sheet_id": None, "header": 1, "service_account": "test/service_account.json", - "columns": {'url': 'link', 'status': 'archive status', 'folder': 'destination folder', 'archive': 'archive location', 'date': 'archive date', 'thumbnail': 'thumbnail', 'timestamp': 'upload timestamp', 'title': 'upload title', 'text': 'text content', 'screenshot': 'screenshot', 'hash': 'hash', 'pdq_hash': 'perceptual hashes', 'wacz': 'wacz', 'replaywebpage': 'replaywebpage'}, + "columns": { + "url": "link", + "status": "archive status", + "folder": "destination folder", + "archive": "archive location", + "date": "archive date", + "thumbnail": "thumbnail", + "timestamp": "upload timestamp", + "title": "upload title", + "text": "text content", + "screenshot": "screenshot", + "hash": "hash", + "pdq_hash": "perceptual hashes", + "wacz": "wacz", + "replaywebpage": "replaywebpage", + }, "allow_worksheets": set(), "block_worksheets": set(), "use_sheet_names_in_stored_paths": True, @@ -78,20 +95,21 @@ def fixed_timestamp(): @pytest.fixture def expected_calls(mock_media, fixed_timestamp): """Fixture for the expected cell updates.""" - return [ - (1, 'status', 'my-archiver: success'), - (1, 'archive', 'http://example.com/screenshot.png'), - (1, 'date', '2025-02-01T00:00:00+00:00'), - (1, 'title', 'Example Title'), - (1, 'text', 'Example Content'), - (1, 'timestamp', '2025-01-01T00:00:00+00:00'), - (1, 'hash', 'not-calculated'), + return [ + (1, "status", "my-archiver: success"), + (1, "archive", "http://example.com/screenshot.png"), + (1, "date", "2025-02-01T00:00:00+00:00"), + (1, "title", "Example Title"), + (1, "text", "Example Content"), + (1, "timestamp", "2025-01-01T00:00:00+00:00"), + (1, "hash", "not-calculated"), # (1, 'screenshot', 'http://example.com/screenshot.png'), # (1, 'thumbnail', '=IMAGE("http://example.com/thumbnail.png")'), # (1, 'wacz', 'http://example.com/browsertrix.wacz'), # (1, 'replaywebpage', 'https://replayweb.page/?source=http%3A%2F%2Fexample.com%2Fbrowsertrix.wacz#view=pages&url=') ] + def test_retrieve_gsheet(gsheets_db, metadata, mock_gworksheet): gw, row = gsheets_db._retrieve_gsheet(metadata) assert gw == mock_gworksheet @@ -100,27 +118,34 @@ def test_retrieve_gsheet(gsheets_db, metadata, mock_gworksheet): def test_started(gsheets_db, mock_metadata, mock_gworksheet): gsheets_db.started(mock_metadata) - mock_gworksheet.set_cell.assert_called_once_with(1, 'status', 'Archive in progress') + mock_gworksheet.set_cell.assert_called_once_with(1, "status", "Archive in progress") + def test_failed(gsheets_db, mock_metadata, mock_gworksheet): reason = "Test failure" gsheets_db.failed(mock_metadata, reason) - mock_gworksheet.set_cell.assert_called_once_with(1, 'status', f'Archive failed {reason}') + mock_gworksheet.set_cell.assert_called_once_with(1, "status", f"Archive failed {reason}") def test_aborted(gsheets_db, mock_metadata, mock_gworksheet): gsheets_db.aborted(mock_metadata) - mock_gworksheet.set_cell.assert_called_once_with(1, 'status', '') + mock_gworksheet.set_cell.assert_called_once_with(1, "status", "") def test_done(gsheets_db, metadata, mock_gworksheet, expected_calls, mocker): - mocker.patch("auto_archiver.modules.gsheet_feeder_db.gsheet_feeder_db.get_current_timestamp", return_value='2025-02-01T00:00:00+00:00') + mocker.patch( + "auto_archiver.modules.gsheet_feeder_db.gsheet_feeder_db.get_current_timestamp", + return_value="2025-02-01T00:00:00+00:00", + ) gsheets_db.done(metadata) mock_gworksheet.batch_set_cell.assert_called_once_with(expected_calls) def test_done_cached(gsheets_db, metadata, mock_gworksheet, mocker): - mocker.patch("auto_archiver.modules.gsheet_feeder_db.gsheet_feeder_db.get_current_timestamp", return_value='2025-02-01T00:00:00+00:00') + mocker.patch( + "auto_archiver.modules.gsheet_feeder_db.gsheet_feeder_db.get_current_timestamp", + return_value="2025-02-01T00:00:00+00:00", + ) gsheets_db.done(metadata, cached=True) # Verify the status message includes "[cached]" @@ -131,15 +156,17 @@ def test_done_cached(gsheets_db, metadata, mock_gworksheet, mocker): def test_done_missing_media(gsheets_db, metadata, mock_gworksheet, mocker): # clear media from metadata metadata.media = [] - mocker.patch("auto_archiver.modules.gsheet_feeder_db.gsheet_feeder_db.get_current_timestamp", return_value='2025-02-01T00:00:00+00:00') + mocker.patch( + "auto_archiver.modules.gsheet_feeder_db.gsheet_feeder_db.get_current_timestamp", + return_value="2025-02-01T00:00:00+00:00", + ) gsheets_db.done(metadata) # Verify nothing media-related gets updated call_args = mock_gworksheet.batch_set_cell.call_args[0][0] - media_fields = {'archive', 'screenshot', 'thumbnail', 'wacz', 'replaywebpage'} + media_fields = {"archive", "screenshot", "thumbnail", "wacz", "replaywebpage"} assert all(call[1] not in media_fields for call in call_args) + def test_safe_status_update(gsheets_db, metadata, mock_gworksheet): gsheets_db._safe_status_update(metadata, "Test status") - mock_gworksheet.set_cell.assert_called_once_with(1, 'status', 'Test status') - - + mock_gworksheet.set_cell.assert_called_once_with(1, "status", "Test status") diff --git a/tests/enrichers/test_hash_enricher.py b/tests/enrichers/test_hash_enricher.py index c2fe67a..05b3b3c 100644 --- a/tests/enrichers/test_hash_enricher.py +++ b/tests/enrichers/test_hash_enricher.py @@ -4,34 +4,50 @@ from auto_archiver.modules.hash_enricher import HashEnricher from auto_archiver.core import Metadata, Media from auto_archiver.core.module import ModuleFactory -@pytest.mark.parametrize("algorithm, filename, expected_hash", [ - ("SHA-256", "tests/data/testfile_1.txt", "1b4f0e9851971998e732078544c96b36c3d01cedf7caa332359d6f1d83567014"), - ("SHA-256", "tests/data/testfile_2.txt", "60303ae22b998861bce3b28f33eec1be758a213c86c93c076dbe9f558c11c752"), - ("SHA3-512", "tests/data/testfile_1.txt", "d2d8cc4f369b340130bd2b29b8b54e918b7c260c3279176da9ccaa37c96eb71735fc97568e892dc6220bf4ae0d748edb46bd75622751556393be3f482e6f794e"), - ("SHA3-512", "tests/data/testfile_2.txt", "e35970edaa1e0d8af7d948491b2da0450a49fd9cc1e83c5db4c6f175f9550cf341f642f6be8cfb0bfa476e4258e5088c5ad549087bf02811132ac2fa22b734c6") -]) + +@pytest.mark.parametrize( + "algorithm, filename, expected_hash", + [ + ("SHA-256", "tests/data/testfile_1.txt", "1b4f0e9851971998e732078544c96b36c3d01cedf7caa332359d6f1d83567014"), + ("SHA-256", "tests/data/testfile_2.txt", "60303ae22b998861bce3b28f33eec1be758a213c86c93c076dbe9f558c11c752"), + ( + "SHA3-512", + "tests/data/testfile_1.txt", + "d2d8cc4f369b340130bd2b29b8b54e918b7c260c3279176da9ccaa37c96eb71735fc97568e892dc6220bf4ae0d748edb46bd75622751556393be3f482e6f794e", + ), + ( + "SHA3-512", + "tests/data/testfile_2.txt", + "e35970edaa1e0d8af7d948491b2da0450a49fd9cc1e83c5db4c6f175f9550cf341f642f6be8cfb0bfa476e4258e5088c5ad549087bf02811132ac2fa22b734c6", + ), + ], +) def test_calculate_hash(algorithm, filename, expected_hash, setup_module): # test SHA-256 he = setup_module(HashEnricher, {"algorithm": algorithm, "chunksize": 100}) assert he.calculate_hash(filename) == expected_hash + def test_default_config_values(setup_module): he = setup_module(HashEnricher) assert he.algorithm == "SHA-256" assert he.chunksize == 16000000 + def test_config(): # test default config - c = ModuleFactory().get_module_lazy('hash_enricher').configs + c = ModuleFactory().get_module_lazy("hash_enricher").configs assert c["algorithm"]["default"] == "SHA-256" assert c["chunksize"]["default"] == 16000000 assert c["algorithm"]["choices"] == ["SHA-256", "SHA3-512"] assert c["algorithm"]["help"] == "hash algorithm to use" - assert c["chunksize"]["help"] == "number of bytes to use when reading files in chunks (if this value is too large you will run out of RAM), default is 16MB" + assert ( + c["chunksize"]["help"] + == "number of bytes to use when reading files in chunks (if this value is too large you will run out of RAM), default is 16MB" + ) def test_hash_media(setup_module): - he = setup_module(HashEnricher, {"algorithm": "SHA-256", "chunksize": 1}) # generate metadata with two test files @@ -46,4 +62,4 @@ def test_hash_media(setup_module): he.enrich(m) assert m.media[0].get("hash") == "SHA-256:1b4f0e9851971998e732078544c96b36c3d01cedf7caa332359d6f1d83567014" - assert m.media[1].get("hash") == "SHA-256:60303ae22b998861bce3b28f33eec1be758a213c86c93c076dbe9f558c11c752" \ No newline at end of file + assert m.media[1].get("hash") == "SHA-256:60303ae22b998861bce3b28f33eec1be758a213c86c93c076dbe9f558c11c752" diff --git a/tests/enrichers/test_meta_enricher.py b/tests/enrichers/test_meta_enricher.py index 476e25b..fe0d737 100644 --- a/tests/enrichers/test_meta_enricher.py +++ b/tests/enrichers/test_meta_enricher.py @@ -1,4 +1,3 @@ -import datetime from datetime import datetime, timedelta, timezone import pytest @@ -16,6 +15,7 @@ def mock_metadata(mocker): mock.get_all_media.return_value = [] return mock + @pytest.fixture def mock_media(mocker): """Creates a mock Media object.""" @@ -59,6 +59,7 @@ def test_enrich_file_sizes(meta_enricher, metadata, tmp_path): assert metadata.get("total_bytes") == 3000 assert metadata.get("total_size") == "2.9 KB" + @pytest.mark.parametrize( "size, expected", [ @@ -74,6 +75,7 @@ def test_human_readable_bytes(size, expected): enricher = MetaEnricher() assert enricher.human_readable_bytes(size) == expected + def test_enrich_file_sizes_no_media(meta_enricher, metadata): """Test that enrich_file_sizes() handles empty media list gracefully.""" meta_enricher.enrich_file_sizes(metadata) @@ -91,4 +93,4 @@ def test_enrich_archive_duration(meta_enricher, metadata, mocker): mock_datetime.now.return_value = mock_now meta_enricher.enrich_archive_duration(metadata) - assert metadata.get("archive_duration_seconds") == 630 \ No newline at end of file + assert metadata.get("archive_duration_seconds") == 630 diff --git a/tests/enrichers/test_metadata_enricher.py b/tests/enrichers/test_metadata_enricher.py index 888837d..14cfc44 100644 --- a/tests/enrichers/test_metadata_enricher.py +++ b/tests/enrichers/test_metadata_enricher.py @@ -1,4 +1,3 @@ - import pytest from auto_archiver.core import Media @@ -33,9 +32,7 @@ def test_get_metadata(enricher, output, expected, mocker): result = enricher.get_metadata("test.jpg") assert result == expected - mock_run.assert_called_once_with( - ["exiftool", "test.jpg"], capture_output=True, text=True - ) + mock_run.assert_called_once_with(["exiftool", "test.jpg"], capture_output=True, text=True) def test_get_metadata_exiftool_not_found(enricher, mocker): @@ -85,4 +82,3 @@ def test_metadata_pickle(enricher, unpickle, mocker): actual_media = metadata.media assert len(expected_media) == len(actual_media) assert actual_media[0].properties.get("metadata") == expected_media[0].properties.get("metadata") - diff --git a/tests/enrichers/test_pdq_hash_enricher.py b/tests/enrichers/test_pdq_hash_enricher.py index 61fa778..483392d 100644 --- a/tests/enrichers/test_pdq_hash_enricher.py +++ b/tests/enrichers/test_pdq_hash_enricher.py @@ -57,7 +57,7 @@ def test_enrich_handles_corrupted_image(metadata_with_images, mocker): ("screenshot", False), ("warc-file-123", False), ("regular-image", True), - ] + ], ) def test_enrich_excludes_by_filetype(media_id, should_have_hash, mocker): metadata = Metadata() @@ -73,4 +73,3 @@ def test_enrich_excludes_by_filetype(media_id, should_have_hash, mocker): media_item = metadata.media[0] assert (media_item.get("pdq_hash") is not None) == should_have_hash - diff --git a/tests/enrichers/test_screenshot_enricher.py b/tests/enrichers/test_screenshot_enricher.py index 25ca51d..b86bb17 100644 --- a/tests/enrichers/test_screenshot_enricher.py +++ b/tests/enrichers/test_screenshot_enricher.py @@ -15,13 +15,15 @@ def mock_selenium_env(mocker): mock_which = mocker.patch("shutil.which") mock_driver_class = mocker.patch("auto_archiver.utils.webdriver.CookieSettingDriver") mock_binary_paths = mocker.patch("selenium.webdriver.common.selenium_manager.SeleniumManager.binary_paths") - mock_is_file = mocker.patch("pathlib.Path.is_file", return_value=True) + mocker.patch("pathlib.Path.is_file", return_value=True) mock_popen = mocker.patch("subprocess.Popen") - mock_is_connectable = mocker.patch("selenium.webdriver.common.service.Service.is_connectable", return_value=True) + mocker.patch("selenium.webdriver.common.service.Service.is_connectable", return_value=True) mock_firefox_options = mocker.patch("selenium.webdriver.FirefoxOptions") + # Define side effect for `shutil.which` def mock_which_side_effect(dep): return "/mock/geckodriver" if dep == "geckodriver" else None + mock_which.side_effect = mock_which_side_effect # Mock binary paths @@ -104,13 +106,7 @@ def test_enrich_adds_screenshot( ], ) def test_enrich_auth_wall( - screenshot_enricher, - metadata_with_video, - mock_selenium_env, - common_patches, - url, - is_auth, - mocker + screenshot_enricher, metadata_with_video, mock_selenium_env, common_patches, url, is_auth, mocker ): # Testing with and without is_auth_wall mock_driver, mock_driver_class, _ = mock_selenium_env @@ -128,9 +124,7 @@ def test_enrich_auth_wall( assert metadata_with_video.media[1].properties.get("id") == "screenshot" -def test_handle_timeout_exception( - screenshot_enricher, metadata_with_video, mock_selenium_env, mocker -): +def test_handle_timeout_exception(screenshot_enricher, metadata_with_video, mock_selenium_env, mocker): mock_driver, mock_driver_class, mock_options_instance = mock_selenium_env mock_driver.get.side_effect = TimeoutException @@ -140,9 +134,7 @@ def test_handle_timeout_exception( assert len(metadata_with_video.media) == 1 -def test_handle_general_exception( - screenshot_enricher, metadata_with_video, mock_selenium_env, mocker -): +def test_handle_general_exception(screenshot_enricher, metadata_with_video, mock_selenium_env, mocker): """Test proper handling of unexpected general exceptions""" mock_driver, mock_driver_class, mock_options_instance = mock_selenium_env # Simulate a generic exception when save_screenshot is called @@ -152,9 +144,7 @@ def test_handle_general_exception( mock_log = mocker.patch("loguru.logger.error") screenshot_enricher.enrich(metadata_with_video) # Verify that the exception was logged with the log - mock_log.assert_called_once_with( - "Got error while loading webdriver for screenshot enricher: Unexpected Error" - ) + mock_log.assert_called_once_with("Got error while loading webdriver for screenshot enricher: Unexpected Error") # And no new media was added due to the error assert len(metadata_with_video.media) == 1 @@ -167,13 +157,12 @@ def test_pdf_creation(mocker, screenshot_enricher, metadata_with_video, mock_sel # Mock the print_page method to return base64-encoded content mock_driver.print_page.return_value = base64.b64encode(b"fake_pdf_content").decode("utf-8") # Patch functions with mocker - mock_os_path_join = mocker.patch("os.path.join", side_effect=lambda *args: f"{args[-1]}") - mock_random_str = mocker.patch( + mocker.patch("os.path.join", side_effect=lambda *args: f"{args[-1]}") + mocker.patch( "auto_archiver.modules.screenshot_enricher.screenshot_enricher.random_str", return_value="fixed123", ) mock_open = mocker.patch("builtins.open", new_callable=mocker.mock_open) - mock_log_error = mocker.patch("loguru.logger.error") screenshot_enricher.enrich(metadata_with_video) # Verify screenshot and PDF creation diff --git a/tests/enrichers/test_ssl_enricher.py b/tests/enrichers/test_ssl_enricher.py index eb7ba6b..cd118ad 100644 --- a/tests/enrichers/test_ssl_enricher.py +++ b/tests/enrichers/test_ssl_enricher.py @@ -51,4 +51,3 @@ def test_ssl_error_handling(enricher, metadata, mocker): mocker.patch("ssl.get_server_certificate", side_effect=ssl.SSLError("SSL error")) with pytest.raises(ssl.SSLError, match="SSL error"): enricher.enrich(metadata) - diff --git a/tests/enrichers/test_thumbnail_enricher.py b/tests/enrichers/test_thumbnail_enricher.py index effc25e..fdc28b7 100644 --- a/tests/enrichers/test_thumbnail_enricher.py +++ b/tests/enrichers/test_thumbnail_enricher.py @@ -25,7 +25,7 @@ def mock_ffmpeg_environment(mocker): # Mocking all the ffmpeg calls in one place mock_ffmpeg_input = mocker.patch("ffmpeg.input") mock_makedirs = mocker.patch("os.makedirs") - mocker.patch.object(Media, "is_video", return_value=True), + (mocker.patch.object(Media, "is_video", return_value=True),) mock_probe = mocker.patch( "ffmpeg.probe", return_value={ @@ -35,9 +35,7 @@ def mock_ffmpeg_environment(mocker): }, ) mock_output = mocker.MagicMock() - mock_ffmpeg_input.return_value.filter.return_value.output.return_value = ( - mock_output - ) + mock_ffmpeg_input.return_value.filter.return_value.output.return_value = mock_output return { "mock_ffmpeg_input": mock_ffmpeg_input, @@ -47,14 +45,21 @@ def mock_ffmpeg_environment(mocker): } -@pytest.mark.parametrize("thumbnails_per_minute, max_thumbnails, expected_count", [ - (10, 5, 5), # Capped at max_thumbnails - (1, 10, 2), # Less than max_thumbnails - (60, 7, 7), # Matches exactly -]) +@pytest.mark.parametrize( + "thumbnails_per_minute, max_thumbnails, expected_count", + [ + (10, 5, 5), # Capped at max_thumbnails + (1, 10, 2), # Less than max_thumbnails + (60, 7, 7), # Matches exactly + ], +) def test_enrich_thumbnail_limits( - thumbnail_enricher, metadata_with_video, mock_ffmpeg_environment, - thumbnails_per_minute, max_thumbnails, expected_count + thumbnail_enricher, + metadata_with_video, + mock_ffmpeg_environment, + thumbnails_per_minute, + max_thumbnails, + expected_count, ): thumbnail_enricher.thumbnails_per_minute = thumbnails_per_minute thumbnail_enricher.max_thumbnails = max_thumbnails @@ -65,8 +70,8 @@ def test_enrich_thumbnail_limits( thumbnails = metadata_with_video.media[0].get("thumbnails") assert len(thumbnails) == expected_count -def test_enrich_handles_probe_failure(thumbnail_enricher, metadata_with_video, mocker): +def test_enrich_handles_probe_failure(thumbnail_enricher, metadata_with_video, mocker): mocker.patch("ffmpeg.probe", side_effect=Exception("Probe error")) mocker.patch("os.makedirs") mock_logger = mocker.patch("loguru.logger.error") @@ -74,36 +79,43 @@ def test_enrich_handles_probe_failure(thumbnail_enricher, metadata_with_video, m thumbnail_enricher.enrich(metadata_with_video) # Ensure error was logged - mock_logger.assert_called_with( - f"error getting duration of video video.mp4: Probe error" - ) + mock_logger.assert_called_with("error getting duration of video video.mp4: Probe error") # Ensure no thumbnails were created thumbnails = metadata_with_video.media[0].get("thumbnails") assert thumbnails is None def test_enrich_skips_non_video_files(thumbnail_enricher, metadata_with_video, mocker): - mocker.patch.object(Media, "is_video", return_value=False) - mock_ffmpeg = mocker.patch("ffmpeg.input") - thumbnail_enricher.enrich(metadata_with_video) - mock_ffmpeg.assert_not_called() + mocker.patch.object(Media, "is_video", return_value=False) + mock_ffmpeg = mocker.patch("ffmpeg.input") + thumbnail_enricher.enrich(metadata_with_video) + mock_ffmpeg.assert_not_called() -@pytest.mark.parametrize("thumbnails_per_minute,max_thumbnails,expected_count", [ - (60, 5, 5), # caught by max - (60, 20, 10), # caught by t/min - (0, 20, 1), # test min caught (1) - (11, 20, 1), # test min caught (1) - (12, 20, 2), # test caught by t/min -]) +@pytest.mark.parametrize( + "thumbnails_per_minute,max_thumbnails,expected_count", + [ + (60, 5, 5), # caught by max + (60, 20, 10), # caught by t/min + (0, 20, 1), # test min caught (1) + (11, 20, 1), # test min caught (1) + (12, 20, 2), # test caught by t/min + ], +) def test_enrich_handles_short_video( - thumbnail_enricher, metadata_with_video, mock_ffmpeg_environment, thumbnails_per_minute, max_thumbnails, expected_count, mocker + thumbnail_enricher, + metadata_with_video, + mock_ffmpeg_environment, + thumbnails_per_minute, + max_thumbnails, + expected_count, + mocker, ): # override mock duration fake_duration = 10 mocker.patch( "ffmpeg.probe", - return_value={ "streams": [{"codec_type": "video", "duration": str(fake_duration)}]}, + return_value={"streams": [{"codec_type": "video", "duration": str(fake_duration)}]}, ) thumbnail_enricher.thumbnails_per_minute = thumbnails_per_minute thumbnail_enricher.max_thumbnails = max_thumbnails @@ -114,9 +126,7 @@ def test_enrich_handles_short_video( assert len(thumbnails) == expected_count -def test_uses_existing_duration( - thumbnail_enricher, metadata_with_video, mock_ffmpeg_environment -): +def test_uses_existing_duration(thumbnail_enricher, metadata_with_video, mock_ffmpeg_environment): metadata_with_video.media[0].set("duration", 60) thumbnail_enricher.enrich(metadata_with_video) mock_ffmpeg_environment["mock_probe"].assert_not_called() @@ -125,7 +135,7 @@ def test_uses_existing_duration( def test_enrich_metadata_structure(thumbnail_enricher, metadata_with_video, mock_ffmpeg_environment, mocker): fake_duration = 120 - mocker.patch("ffmpeg.probe", return_value={'streams': [{'codec_type': 'video', 'duration': str(fake_duration)}]}) + mocker.patch("ffmpeg.probe", return_value={"streams": [{"codec_type": "video", "duration": str(fake_duration)}]}) thumbnail_enricher.thumbnails_per_minute = 2 thumbnail_enricher.max_thumbnails = 4 diff --git a/tests/enrichers/test_wayback_enricher.py b/tests/enrichers/test_wayback_enricher.py index 5406e39..113458b 100644 --- a/tests/enrichers/test_wayback_enricher.py +++ b/tests/enrichers/test_wayback_enricher.py @@ -5,37 +5,52 @@ from auto_archiver.modules.wayback_extractor_enricher import WaybackExtractorEnr from auto_archiver.core import Metadata +@pytest.fixture(autouse=True) +def mock_sleep(mocker): + """Mock time.sleep to avoid delays.""" + return mocker.patch("time.sleep") + + @pytest.fixture def mock_is_auth_wall(mocker): """Fixture to mock is_auth_wall behavior.""" + def _mock_is_auth_wall(return_value: bool): return mocker.patch("auto_archiver.utils.url.is_auth_wall", return_value=return_value) + return _mock_is_auth_wall + @pytest.fixture def mock_post_success(mocker): """Fixture to mock POST requests with a successful response.""" + def _mock_post(json_data: dict = None, status_code: int = 200): json_data = {"job_id": "job123"} if json_data is None else json_data resp = mocker.Mock(status_code=status_code) resp.json.return_value = json_data return mocker.patch("requests.post", return_value=resp) + return _mock_post + @pytest.fixture def mock_get_success(mocker): """Fixture to mock GET requests returning a completed archive status.""" + def _mock_get(json_data: dict = None, status_code: int = 200): json_data = json_data or { "status": "success", "timestamp": "20250101010101", - "original_url": "https://example.com" + "original_url": "https://example.com", } resp = mocker.Mock(status_code=status_code) resp.json.return_value = json_data return mocker.patch("requests.get", return_value=resp) + return _mock_get + @pytest.fixture def wayback_extractor_enricher(setup_module) -> WaybackExtractorEnricher: configs: dict = { @@ -49,12 +64,7 @@ def wayback_extractor_enricher(setup_module) -> WaybackExtractorEnricher: return setup_module("wayback_extractor_enricher", configs) -def test_download_success( - wayback_extractor_enricher, - mock_is_auth_wall, - mock_post_success, - mock_get_success -): +def test_download_success(wayback_extractor_enricher, mock_is_auth_wall, mock_post_success, mock_get_success): mock_is_auth_wall(False) mock_post_success() mock_get_success() @@ -63,34 +73,28 @@ def test_download_success( result = wayback_extractor_enricher.download(metadata) assert result.get("wayback") == "https://web.archive.org/web/20250101010101/https://example.com" + def test_enrich_auth_wall(wayback_extractor_enricher, metadata, mock_is_auth_wall): mock_is_auth_wall(True) result = wayback_extractor_enricher.enrich(metadata) assert result is None + def test_enrich_already_enriched(wayback_extractor_enricher, metadata): metadata.set("wayback", "existing") result = wayback_extractor_enricher.enrich(metadata) assert result is True -def test_enrich_post_failure( - wayback_extractor_enricher, - metadata, - mock_is_auth_wall, - mock_post_success -): + +def test_enrich_post_failure(wayback_extractor_enricher, metadata, mock_is_auth_wall, mock_post_success): mock_is_auth_wall(False) mock_post_success(json_data={"error": "server error"}, status_code=500) result = wayback_extractor_enricher.enrich(metadata) assert result is False assert "Internet archive failed with status of 500" in metadata.get("wayback") -def test_enrich_post_json_decode_error( - wayback_extractor_enricher, - metadata, - mock_is_auth_wall, - mocker -): + +def test_enrich_post_json_decode_error(wayback_extractor_enricher, metadata, mock_is_auth_wall, mocker): mock_is_auth_wall(False) resp = mocker.Mock(status_code=200) resp.json.side_effect = json.decoder.JSONDecodeError("msg", "doc", 0) @@ -98,22 +102,15 @@ def test_enrich_post_json_decode_error( mocker.patch("requests.post", return_value=resp) assert wayback_extractor_enricher.enrich(metadata) is False -def test_enrich_no_job_id( - wayback_extractor_enricher, - metadata, - mock_is_auth_wall, - mock_post_success -): + +def test_enrich_no_job_id(wayback_extractor_enricher, metadata, mock_is_auth_wall, mock_post_success): mock_is_auth_wall(False) mock_post_success(json_data={}) assert wayback_extractor_enricher.enrich(metadata) is False + def test_enrich_get_success( - wayback_extractor_enricher, - metadata, - mock_is_auth_wall, - mock_post_success, - mock_get_success + wayback_extractor_enricher, metadata, mock_is_auth_wall, mock_post_success, mock_get_success ): mock_is_auth_wall(False) mock_post_success() @@ -122,24 +119,18 @@ def test_enrich_get_success( assert metadata.get("wayback") == "https://web.archive.org/web/20250101010101/https://example.com" assert metadata.get("check wayback") == "https://web.archive.org/web/*/https://example.com" + def test_enrich_get_failure( - wayback_extractor_enricher, - metadata, - mock_is_auth_wall, - mock_post_success, - mock_get_success + wayback_extractor_enricher, metadata, mock_is_auth_wall, mock_post_success, mock_get_success ): mock_is_auth_wall(False) mock_post_success() mock_get_success(json_data={"status": "failed"}, status_code=400) assert wayback_extractor_enricher.enrich(metadata) is False + def test_enrich_get_request_exception( - wayback_extractor_enricher, - metadata, - mock_is_auth_wall, - mock_post_success, - mocker + wayback_extractor_enricher, metadata, mock_is_auth_wall, mock_post_success, mocker ): mock_is_auth_wall(False) mock_post_success() @@ -149,12 +140,9 @@ def test_enrich_get_request_exception( assert wayback_extractor_enricher.enrich(metadata) is True assert metadata.get("wayback").get("job_id") == "job123" + def test_enrich_get_json_decode_error( - wayback_extractor_enricher, - metadata, - mock_is_auth_wall, - mock_post_success, - mocker + wayback_extractor_enricher, metadata, mock_is_auth_wall, mock_post_success, mocker ): mock_is_auth_wall(False) mock_post_success() diff --git a/tests/enrichers/test_whisper_enricher.py b/tests/enrichers/test_whisper_enricher.py index ee1844a..f669631 100644 --- a/tests/enrichers/test_whisper_enricher.py +++ b/tests/enrichers/test_whisper_enricher.py @@ -7,6 +7,12 @@ from auto_archiver.modules.whisper_enricher import WhisperEnricher TEST_S3_URL = "http://cdn.example.com/test.mp4" +@pytest.fixture(autouse=True) +def mock_sleep(mocker): + """Mock time.sleep to avoid delays.""" + return mocker.patch("time.sleep") + + @pytest.fixture def enricher(mocker): """Fixture with mocked S3 and API dependencies""" @@ -16,7 +22,7 @@ def enricher(mocker): "include_srt": False, "timeout": 5, "action": "translate", - "steps": {"storages": ["s3_storage"]} + "steps": {"storages": ["s3_storage"]}, } mock_s3 = mocker.MagicMock(spec=S3Storage) mock_s3.get_cdn_url.return_value = TEST_S3_URL @@ -25,7 +31,7 @@ def enricher(mocker): instance.display_name = "Whisper Enricher" instance.config_setup({instance.name: config}) # bypassing the setup method and mocking S3 setup - instance.stores = config['steps']['storages'] + instance.stores = config["steps"]["storages"] instance.s3 = mock_s3 yield instance, mock_s3 @@ -63,19 +69,14 @@ def test_successful_job_submission(enricher, metadata, mock_requests, mocker): # Mock the complete API interaction chain mock_status_response = mocker.MagicMock() mock_status_response.status_code = 200 - mock_status_response.json.return_value = { - "status": "success", - "meta": {} - } + mock_status_response.json.return_value = {"status": "success", "meta": {}} mock_artifacts_response = mocker.MagicMock() mock_artifacts_response.status_code = 200 - mock_artifacts_response.json.return_value = [{ - "data": [{"start": 0, "end": 5, "text": "test transcript"}] - }] + mock_artifacts_response.json.return_value = [{"data": [{"start": 0, "end": 5, "text": "test transcript"}]}] # Set up mock response sequence mock_requests.get.side_effect = [ mock_status_response, # First call: status check - mock_artifacts_response # Second call: artifacts check + mock_artifacts_response, # Second call: artifacts check ] # Run enrichment (without opening file) @@ -84,15 +85,17 @@ def test_successful_job_submission(enricher, metadata, mock_requests, mocker): mock_requests.post.assert_called_once_with( "http://testapi/jobs", json={"url": "http://cdn.example.com/test.mp4", "type": "translate"}, - headers={"Authorization": "Bearer whisper-key"} + headers={"Authorization": "Bearer whisper-key"}, ) # Verify job status checks assert mock_requests.get.call_count == 2 assert "artifact_0_text" in metadata.media[0].get("whisper_model") - assert metadata.media[0].get("whisper_model") == {'artifact_0_text': 'test transcript', - 'job_artifacts_check': 'http://testapi/jobs/job123/artifacts', - 'job_id': 'job123', - 'job_status_check': 'http://testapi/jobs/job123'} + assert metadata.media[0].get("whisper_model") == { + "artifact_0_text": "test transcript", + "job_artifacts_check": "http://testapi/jobs/job123/artifacts", + "job_id": "job123", + "job_status_check": "http://testapi/jobs/job123", + } def test_submit_job(enricher, mocker): diff --git a/tests/extractors/test_extractor_base.py b/tests/extractors/test_extractor_base.py index 6e77ec3..0240529 100644 --- a/tests/extractors/test_extractor_base.py +++ b/tests/extractors/test_extractor_base.py @@ -7,7 +7,6 @@ from auto_archiver.core.extractor import Extractor class TestExtractorBase(object): - extractor_module: str = None config: dict = None @@ -17,7 +16,7 @@ class TestExtractorBase(object): assert self.config is not None, "self.config must be a dict set on the subclass" self.extractor: Type[Extractor] = setup_module(self.extractor_module, self.config) - + def assertValidResponseMetadata(self, test_response: Metadata, title: str, timestamp: str, status: str = ""): assert test_response is not False diff --git a/tests/extractors/test_generic_extractor.py b/tests/extractors/test_generic_extractor.py index 33f35b7..ec7aae4 100644 --- a/tests/extractors/test_generic_extractor.py +++ b/tests/extractors/test_generic_extractor.py @@ -9,26 +9,28 @@ import pytest from auto_archiver.modules.generic_extractor.generic_extractor import GenericExtractor from .test_extractor_base import TestExtractorBase -CI=os.getenv("GITHUB_ACTIONS", '') == 'true' +CI = os.getenv("GITHUB_ACTIONS", "") == "true" + + class TestGenericExtractor(TestExtractorBase): - """Tests Generic Extractor - """ - extractor_module = 'generic_extractor' + """Tests Generic Extractor""" + + extractor_module = "generic_extractor" extractor: GenericExtractor config = { - 'subtitles': False, - 'comments': False, - 'livestreams': False, - 'live_from_start': False, - 'end_means_success': True, - 'allow_playlist': False, - 'max_downloads': "inf", - 'proxy': None, - 'cookies_from_browser': False, - 'cookie_file': None, - } - + "subtitles": False, + "comments": False, + "livestreams": False, + "live_from_start": False, + "end_means_success": True, + "allow_playlist": False, + "max_downloads": "inf", + "proxy": None, + "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.modules.generic_extractor" @@ -38,21 +40,26 @@ class TestGenericExtractor(TestExtractorBase): path = os.path.join(dirname(dirname(__file__)), "data/") assert self.extractor.dropin_for_name("dropin", additional_paths=[path]) - - - @pytest.mark.parametrize("url, is_suitable", [ - ("https://www.youtube.com/watch?v=5qap5aO4i9A", True), - ("https://www.tiktok.com/@funnycats0ftiktok/video/7345101300750748970?lang=en", True), - ("https://www.instagram.com/p/CU1J9JYJ9Zz/", True), - ("https://www.facebook.com/nytimes/videos/10160796550110716", True), - ("https://www.twitch.tv/videos/1167226570", True), - ("https://bellingcat.com/news/2021/10/08/ukrainian-soldiers-are-being-killed-by-landmines-in-the-donbas/", True), - ("https://google.com", True)]) + @pytest.mark.parametrize( + "url, is_suitable", + [ + ("https://www.youtube.com/watch?v=5qap5aO4i9A", True), + ("https://www.tiktok.com/@funnycats0ftiktok/video/7345101300750748970?lang=en", True), + ("https://www.instagram.com/p/CU1J9JYJ9Zz/", True), + ("https://www.facebook.com/nytimes/videos/10160796550110716", True), + ("https://www.twitch.tv/videos/1167226570", True), + ( + "https://bellingcat.com/news/2021/10/08/ukrainian-soldiers-are-being-killed-by-landmines-in-the-donbas/", + True, + ), + ("https://google.com", True), + ], + ) def test_suitable_urls(self, make_item, url, is_suitable): """ - Note: expected behaviour is to return True for all URLs, as YoutubeDLArchiver should be able to handle all URLs - This behaviour may be changed in the future (e.g. if we want the youtubedl archiver to just handle URLs it has extractors for, - and then if and only if all archivers fails, does it fall back to the generic archiver) + Note: expected behaviour is to return True for all URLs, as YoutubeDLArchiver should be able to handle all URLs + This behaviour may be changed in the future (e.g. if we want the youtubedl archiver to just handle URLs it has extractors for, + and then if and only if all archivers fails, does it fall back to the generic archiver) """ assert self.extractor.suitable(url) == is_suitable @@ -63,11 +70,14 @@ class TestGenericExtractor(TestExtractorBase): assert result.get_url() == "https://www.tiktok.com/@funnycats0ftiktok/video/7345101300750748970" @pytest.mark.download - @pytest.mark.parametrize("url", [ - "https://bsky.app/profile/colborne.bsky.social/post/3lcxcpgt6j42l", - "twitter.com/bellingcat/status/123", - "https://www.youtube.com/watch?v=1" - ]) + @pytest.mark.parametrize( + "url", + [ + "https://bsky.app/profile/colborne.bsky.social/post/3lcxcpgt6j42l", + "twitter.com/bellingcat/status/123", + "https://www.youtube.com/watch?v=1", + ], + ) def test_download_nonexistent_media(self, make_item, url): """ Test to make sure that the extractor doesn't break on non-existend posts/media @@ -78,7 +88,10 @@ class TestGenericExtractor(TestExtractorBase): result = self.extractor.download(item) assert not result - @pytest.mark.skipif(CI, reason="Currently no way to authenticate when on CI. Youtube (yt-dlp) doesn't support logging in with username/password.") + @pytest.mark.skipif( + CI, + reason="Currently no way to authenticate when on CI. Youtube (yt-dlp) doesn't support logging in with username/password.", + ) @pytest.mark.download def test_youtube_download(self, make_item): # url https://www.youtube.com/watch?v=5qap5aO4i9A @@ -87,7 +100,10 @@ class TestGenericExtractor(TestExtractorBase): result = self.extractor.download(item) assert result.get_url() == "https://www.youtube.com/watch?v=J---aiyznGQ" assert result.get_title() == "Keyboard Cat! - THE ORIGINAL!" - assert result.get('description') == "Buy NEW Keyboard Cat Merch! https://keyboardcat.creator-spring.com\n\nxo Keyboard Cat memes make your day better!\nhttp://www.keyboardcatstore.com/\nhttps://www.facebook.com/thekeyboardcat\nhttp://www.charlieschmidt.com/" + assert ( + result.get("description") + == "Buy NEW Keyboard Cat Merch! https://keyboardcat.creator-spring.com\n\nxo Keyboard Cat memes make your day better!\nhttp://www.keyboardcatstore.com/\nhttps://www.facebook.com/thekeyboardcat\nhttp://www.charlieschmidt.com/" + ) assert len(result.media) == 2 assert Path(result.media[0].filename).name == "J---aiyznGQ.webm" assert Path(result.media[1].filename).name == "hqdefault.jpg" @@ -103,7 +119,7 @@ class TestGenericExtractor(TestExtractorBase): item = make_item("https://bsky.app/profile/bellingcat.com/post/3lfn3hbcxgc2q") result = self.extractor.download(item) assert result is not False - + @pytest.mark.download def test_bluesky_download_no_media(self, make_item): item = make_item("https://bsky.app/profile/bellingcat.com/post/3lfphwmcs4c2z") @@ -115,7 +131,7 @@ class TestGenericExtractor(TestExtractorBase): item = make_item("https://bsky.app/profile/bellingcat.com/post/3le2l4gsxlk2i") result = self.extractor.download(item) assert result is not False - + @pytest.mark.skipif(CI, reason="Truth social blocks GH actions.") @pytest.mark.download def test_truthsocial_download_video(self, make_item): @@ -130,14 +146,14 @@ class TestGenericExtractor(TestExtractorBase): item = make_item("https://truthsocial.com/@bbcnewa/posts/109598702184774628") result = self.extractor.download(item) assert result is not False - + @pytest.mark.skipif(CI, reason="Truth social blocks GH actions.") @pytest.mark.download def test_truthsocial_download_poll(self, make_item): item = make_item("https://truthsocial.com/@CNN_US/posts/113724326568555098") result = self.extractor.download(item) assert result is not False - + @pytest.mark.skipif(CI, reason="Truth social blocks GH actions.") @pytest.mark.download def test_truthsocial_download_single_image(self, make_item): @@ -159,7 +175,7 @@ class TestGenericExtractor(TestExtractorBase): url = "https://x.com/Bellingcat/status/17197025860711058" response = self.extractor.download(make_item(url)) assert not response - + @pytest.mark.download def test_twitter_download_malformed_tweetid(self, make_item): # this tweet does not exist @@ -169,7 +185,6 @@ class TestGenericExtractor(TestExtractorBase): @pytest.mark.download def test_twitter_download_tweet_no_media(self, make_item): - item = make_item("https://twitter.com/MeCookieMonster/status/1617921633456640001?s=20&t=3d0g4ZQis7dCbSDg-mE7-w") post = self.extractor.download(item) @@ -177,9 +192,9 @@ class TestGenericExtractor(TestExtractorBase): post, "Onion rings are just vegetable donuts.", datetime.datetime(2023, 1, 24, 16, 25, 51, tzinfo=datetime.timezone.utc), - "yt-dlp_Twitter: success" + "yt-dlp_Twitter: success", ) - + @pytest.mark.download def test_twitter_download_video(self, make_item): url = "https://x.com/bellingcat/status/1871552600346415571" @@ -187,26 +202,46 @@ class TestGenericExtractor(TestExtractorBase): self.assertValidResponseMetadata( post, "Bellingcat - This month's Bellingchat Premium is with @KolinaKoltai. She reveals how she investigated a platform allowing users to create AI-generated child sexual abuse material and explains why it's crucial to investigate the people behind these services", - datetime.datetime(2024, 12, 24, 13, 44, 46, tzinfo=datetime.timezone.utc) + datetime.datetime(2024, 12, 24, 13, 44, 46, tzinfo=datetime.timezone.utc), ) - @pytest.mark.xfail(reason="Currently failing, sensitive content requires logged in users/cookies - not yet implemented") + @pytest.mark.xfail( + reason="Currently failing, sensitive content requires logged in users/cookies - not yet implemented" + ) @pytest.mark.download - @pytest.mark.parametrize("url, title, timestamp, image_hash", [ - ("https://x.com/SozinhoRamalho/status/1876710769913450647", "ignore tweet, testing sensitivity warning nudity", datetime.datetime(2024, 12, 31, 14, 18, 33, tzinfo=datetime.timezone.utc), "image_hash"), - ("https://x.com/SozinhoRamalho/status/1876710875475681357", "ignore tweet, testing sensitivity warning violence", datetime.datetime(2024, 12, 31, 14, 18, 33, tzinfo=datetime.timezone.utc), "image_hash"), - ("https://x.com/SozinhoRamalho/status/1876711053813227618", "ignore tweet, testing sensitivity warning sensitive", datetime.datetime(2024, 12, 31, 14, 18, 33, tzinfo=datetime.timezone.utc), "image_hash"), - ("https://x.com/SozinhoRamalho/status/1876711141314801937", "ignore tweet, testing sensitivity warning nudity, violence, sensitivity", datetime.datetime(2024, 12, 31, 14, 18, 33, tzinfo=datetime.timezone.utc), "image_hash"), - ]) + @pytest.mark.parametrize( + "url, title, timestamp, image_hash", + [ + ( + "https://x.com/SozinhoRamalho/status/1876710769913450647", + "ignore tweet, testing sensitivity warning nudity", + datetime.datetime(2024, 12, 31, 14, 18, 33, tzinfo=datetime.timezone.utc), + "image_hash", + ), + ( + "https://x.com/SozinhoRamalho/status/1876710875475681357", + "ignore tweet, testing sensitivity warning violence", + datetime.datetime(2024, 12, 31, 14, 18, 33, tzinfo=datetime.timezone.utc), + "image_hash", + ), + ( + "https://x.com/SozinhoRamalho/status/1876711053813227618", + "ignore tweet, testing sensitivity warning sensitive", + datetime.datetime(2024, 12, 31, 14, 18, 33, tzinfo=datetime.timezone.utc), + "image_hash", + ), + ( + "https://x.com/SozinhoRamalho/status/1876711141314801937", + "ignore tweet, testing sensitivity warning nudity, violence, sensitivity", + datetime.datetime(2024, 12, 31, 14, 18, 33, tzinfo=datetime.timezone.utc), + "image_hash", + ), + ], + ) def test_twitter_download_sensitive_media(self, url, title, timestamp, image_hash, make_item): - """Download tweets with sensitive media""" post = self.extractor.download(make_item(url)) - self.assertValidResponseMetadata( - post, - title, - timestamp - ) + self.assertValidResponseMetadata(post, title, timestamp) assert len(post.media) == 1 - assert post.media[0].hash == image_hash \ No newline at end of file + assert post.media[0].hash == image_hash diff --git a/tests/extractors/test_instagram_api_extractor.py b/tests/extractors/test_instagram_api_extractor.py index 7eba8e9..d8a1cc0 100644 --- a/tests/extractors/test_instagram_api_extractor.py +++ b/tests/extractors/test_instagram_api_extractor.py @@ -15,10 +15,11 @@ def mock_user_response(): "username": "test_user", "full_name": "Test User", "profile_pic_url_hd": "http://example.com/profile.jpg", - "profile_pic_url": "http://example.com/profile_lowres.jpg" + "profile_pic_url": "http://example.com/profile_lowres.jpg", } } + @pytest.fixture def mock_post_response(): return { @@ -27,16 +28,14 @@ def mock_post_response(): "caption_text": "Test Caption", "taken_at": datetime.now().timestamp(), "video_url": "http://example.com/video.mp4", - "thumbnail_url": "http://example.com/thumbnail.jpg" + "thumbnail_url": "http://example.com/thumbnail.jpg", } + @pytest.fixture def mock_story_response(): - return [{ - "id": "story_123", - "taken_at": datetime.now().timestamp(), - "video_url": "http://example.com/story.mp4" - }] + return [{"id": "story_123", "taken_at": datetime.now().timestamp(), "video_url": "http://example.com/story.mp4"}] + @pytest.fixture def mock_highlight_response(): @@ -46,11 +45,13 @@ def mock_highlight_response(): "highlight:123": { "id": "123", "title": "Test Highlight", - "items": [{ - "id": "item_123", - "taken_at": datetime.now().timestamp(), - "video_url": "http://example.com/highlight.mp4" - }] + "items": [ + { + "id": "item_123", + "taken_at": datetime.now().timestamp(), + "video_url": "http://example.com/highlight.mp4", + } + ], } } } @@ -81,24 +82,30 @@ class TestInstagramAPIExtractor(TestExtractorBase): m.set("netloc", "instagram.com") return m - @pytest.mark.parametrize("url,expected", [ - ("https://instagram.com/user", [("", "user", "")]), - ("https://instagr.am/p/post_id", []), - ("https://youtube.com", []), - ("https://www.instagram.com/reel/reel_id", [("reel", "reel_id", "")]), - ("https://instagram.com/stories/highlights/123", [("stories/highlights", "123", "")]), - ("https://instagram.com/stories/user/123", [("stories", "user", "123")]), - ]) + @pytest.mark.parametrize( + "url,expected", + [ + ("https://instagram.com/user", [("", "user", "")]), + ("https://instagr.am/p/post_id", []), + ("https://youtube.com", []), + ("https://www.instagram.com/reel/reel_id", [("reel", "reel_id", "")]), + ("https://instagram.com/stories/highlights/123", [("stories/highlights", "123", "")]), + ("https://instagram.com/stories/user/123", [("stories", "user", "123")]), + ], + ) def test_url_parsing(self, url, expected): assert self.extractor.valid_url.findall(url) == expected def test_initialize(self): assert self.extractor.api_endpoint[-1] != "/" - @pytest.mark.parametrize("input_dict,expected", [ - ({"x": 0, "valid": "data"}, {"valid": "data"}), - ({"nested": {"y": None, "valid": [{}]}}, {"nested": {"valid": [{}]}}), - ]) + @pytest.mark.parametrize( + "input_dict,expected", + [ + ({"x": 0, "valid": "data"}, {"valid": "data"}), + ({"nested": {"y": None, "valid": [{}]}}, {"nested": {"valid": [{}]}}), + ], + ) def test_cleanup_dict(self, input_dict, expected): assert self.extractor.cleanup_dict(input_dict) == expected @@ -114,8 +121,8 @@ class TestInstagramAPIExtractor(TestExtractorBase): def test_download_profile_basic(self, metadata, mock_user_response, mocker): """Test basic profile download without full_profile""" - mock_call = mocker.patch.object(self.extractor, 'call_api') - mock_download = mocker.patch.object(self.extractor, 'download_from_url') + mock_call = mocker.patch.object(self.extractor, "call_api") + mock_download = mocker.patch.object(self.extractor, "download_from_url") # Mock API responses mock_call.return_value = mock_user_response mock_download.return_value = "profile.jpg" @@ -132,17 +139,14 @@ class TestInstagramAPIExtractor(TestExtractorBase): def test_download_profile_full(self, metadata, mock_user_response, mock_story_response, mocker): """Test full profile download with stories/posts""" - mock_call = mocker.patch.object(self.extractor, 'call_api') - mock_posts = mocker.patch.object(self.extractor, 'download_all_posts') - mock_highlights = mocker.patch.object(self.extractor, 'download_all_highlights') - mock_tagged = mocker.patch.object(self.extractor, 'download_all_tagged') - mock_stories = mocker.patch.object(self.extractor, '_download_stories_reusable') + mock_call = mocker.patch.object(self.extractor, "call_api") + mock_posts = mocker.patch.object(self.extractor, "download_all_posts") + mock_highlights = mocker.patch.object(self.extractor, "download_all_highlights") + mock_tagged = mocker.patch.object(self.extractor, "download_all_tagged") + mock_stories = mocker.patch.object(self.extractor, "_download_stories_reusable") self.extractor.full_profile = True - mock_call.side_effect = [ - mock_user_response, - mock_story_response - ] + mock_call.side_effect = [mock_user_response, mock_story_response] mock_highlights.return_value = None mock_stories.return_value = mock_story_response mock_posts.return_value = None @@ -155,7 +159,7 @@ class TestInstagramAPIExtractor(TestExtractorBase): def test_download_profile_not_found(self, metadata, mocker): """Test profile not found error""" - mock_call = mocker.patch.object(self.extractor, 'call_api') + mock_call = mocker.patch.object(self.extractor, "call_api") mock_call.return_value = {"user": None} with pytest.raises(AssertionError) as exc_info: self.extractor.download_profile(metadata, "invalid_user") @@ -163,18 +167,14 @@ class TestInstagramAPIExtractor(TestExtractorBase): def test_download_profile_error_handling(self, metadata, mock_user_response, mocker): """Test error handling in full profile mode""" - mock_call = mocker.patch.object(self.extractor, 'call_api') - mock_highlights = mocker.patch.object(self.extractor, 'download_all_highlights') - mock_tagged = mocker.patch.object(self.extractor, 'download_all_tagged') - stories_tagged = mocker.patch.object(self.extractor, '_download_stories_reusable') - mock_posts = mocker.patch.object(self.extractor, 'download_all_posts') + mock_call = mocker.patch.object(self.extractor, "call_api") + mock_highlights = mocker.patch.object(self.extractor, "download_all_highlights") + mock_tagged = mocker.patch.object(self.extractor, "download_all_tagged") + stories_tagged = mocker.patch.object(self.extractor, "_download_stories_reusable") + mock_posts = mocker.patch.object(self.extractor, "download_all_posts") self.extractor.full_profile = True - mock_call.side_effect = [ - mock_user_response, - Exception("Stories API failed"), - Exception("Posts API failed") - ] + mock_call.side_effect = [mock_user_response, Exception("Stories API failed"), Exception("Posts API failed")] mock_highlights.return_value = None mock_tagged.return_value = None stories_tagged.return_value = None @@ -182,4 +182,4 @@ class TestInstagramAPIExtractor(TestExtractorBase): result = self.extractor.download_profile(metadata, "test_user") assert result.is_success() - assert "Error downloading stories for test_user" in result.metadata["errors"] \ No newline at end of file + assert "Error downloading stories for test_user" in result.metadata["errors"] diff --git a/tests/extractors/test_instagram_extractor.py b/tests/extractors/test_instagram_extractor.py index 647cab4..0cafa2b 100644 --- a/tests/extractors/test_instagram_extractor.py +++ b/tests/extractors/test_instagram_extractor.py @@ -5,8 +5,7 @@ from auto_archiver.modules.instagram_extractor import InstagramExtractor @pytest.fixture def instagram_extractor(setup_module, mocker): - - extractor_module: str = 'instagram_extractor' + extractor_module: str = "instagram_extractor" config: dict = { "username": "user_name", "password": "password123", @@ -17,20 +16,26 @@ def instagram_extractor(setup_module, mocker): fake_loader.load_session_from_file.return_value = None fake_loader.login.return_value = None fake_loader.save_session_to_file.return_value = None - mocker.patch("instaloader.Instaloader", return_value=fake_loader,) + mocker.patch( + "instaloader.Instaloader", + return_value=fake_loader, + ) return setup_module(extractor_module, config) -@pytest.mark.parametrize("url", [ - "https://www.instagram.com/p/", - "https://www.instagram.com/p/1234567890/", - "https://www.instagram.com/reel/1234567890/", - "https://www.instagram.com/username/", - "https://www.instagram.com/username/stories/", - "https://www.instagram.com/username/highlights/", -]) +@pytest.mark.parametrize( + "url", + [ + "https://www.instagram.com/p/", + "https://www.instagram.com/p/1234567890/", + "https://www.instagram.com/reel/1234567890/", + "https://www.instagram.com/username/", + "https://www.instagram.com/username/stories/", + "https://www.instagram.com/username/highlights/", + ], +) def test_regex_matches(url: str, instagram_extractor: InstagramExtractor) -> None: """ Ensure that the valid_url regex matches all provided Instagram URLs. """ - assert instagram_extractor.valid_url.match(url) \ No newline at end of file + assert instagram_extractor.valid_url.match(url) diff --git a/tests/extractors/test_instagram_tbot_extractor.py b/tests/extractors/test_instagram_tbot_extractor.py index f274728..47a4bec 100644 --- a/tests/extractors/test_instagram_tbot_extractor.py +++ b/tests/extractors/test_instagram_tbot_extractor.py @@ -7,10 +7,16 @@ from auto_archiver.modules.instagram_tbot_extractor import InstagramTbotExtracto from tests.extractors.test_extractor_base import TestExtractorBase +@pytest.fixture(autouse=True) +def mock_sleep(mocker): + """Mock time.sleep to avoid delays.""" + return mocker.patch("time.sleep") + + @pytest.fixture def patch_extractor_methods(request, setup_module, mocker): - mocker.patch.object(InstagramTbotExtractor, '_prepare_session_file', return_value=None) - mocker.patch.object(InstagramTbotExtractor, '_initialize_telegram_client', return_value=None) + mocker.patch.object(InstagramTbotExtractor, "_prepare_session_file", return_value=None) + mocker.patch.object(InstagramTbotExtractor, "_initialize_telegram_client", return_value=None) yield @@ -35,12 +41,7 @@ def mock_telegram_client(mocker): @pytest.fixture def extractor(setup_module, patch_extractor_methods, mocker): extractor_module = "instagram_tbot_extractor" - config = { - "api_id": 12345, - "api_hash": "test_api_hash", - "session_file": "test_session", - "timeout": 4 - } + config = {"api_id": 12345, "api_hash": "test_api_hash", "session_file": "test_session", "timeout": 4} extractor = setup_module(extractor_module, config) extractor.client = mocker.MagicMock() extractor.session_file = "test_session" @@ -79,21 +80,30 @@ class TestInstagramTbotExtractorReal(TestExtractorBase): "session_file": "secrets/anon-insta", } - @pytest.mark.parametrize("url, expected_status, message, len_media", [ - ("https://www.instagram.com/p/C4QgLbrIKXG", "insta-via-bot: success", - "Are you new to Bellingcat? - The way we share our investigations is different. 💭\nWe want you to read our story but also learn ou", - 6), - ("https://www.instagram.com/reel/DEVLK8qoIbg/", "insta-via-bot: success", - "Our volunteer community is at the centre of many incredible Bellingcat investigations and tools. Stephanie Ladel is one such vol", - 3), - # instagram tbot not working (potentially intermittently?) for stories - replace with a live story to retest - # ("https://www.instagram.com/stories/bellingcatofficial/3556336382743057476/", False, "Media not found or unavailable"), - # Seems to be working intermittently for highlights - # ("https://www.instagram.com/stories/highlights/17868810693068139/", "insta-via-bot: success", None, 50), - # Marking invalid url as success - ("https://www.instagram.com/p/INVALID", "insta-via-bot: success", "Media not found or unavailable", 0), - ("https://www.youtube.com/watch?v=ymCMy8OffHM", False, None, 0), - ]) + @pytest.mark.parametrize( + "url, expected_status, message, len_media", + [ + ( + "https://www.instagram.com/p/C4QgLbrIKXG", + "insta-via-bot: success", + "Are you new to Bellingcat? - The way we share our investigations is different. 💭\nWe want you to read our story but also learn ou", + 6, + ), + ( + "https://www.instagram.com/reel/DEVLK8qoIbg/", + "insta-via-bot: success", + "Our volunteer community is at the centre of many incredible Bellingcat investigations and tools. Stephanie Ladel is one such vol", + 3, + ), + # instagram tbot not working (potentially intermittently?) for stories - replace with a live story to retest + # ("https://www.instagram.com/stories/bellingcatofficial/3556336382743057476/", False, "Media not found or unavailable"), + # Seems to be working intermittently for highlights + # ("https://www.instagram.com/stories/highlights/17868810693068139/", "insta-via-bot: success", None, 50), + # Marking invalid url as success + ("https://www.instagram.com/p/INVALID", "insta-via-bot: success", "Media not found or unavailable", 0), + ("https://www.youtube.com/watch?v=ymCMy8OffHM", False, None, 0), + ], + ) def test_download(self, url, expected_status, message, len_media, metadata_sample): """Test the `download()` method with various Instagram URLs.""" metadata_sample.set_url(url) diff --git a/tests/extractors/test_tiktok_tikwm_extractor.py b/tests/extractors/test_tiktok_tikwm_extractor.py index e8ad8df..d04d7e4 100644 --- a/tests/extractors/test_tiktok_tikwm_extractor.py +++ b/tests/extractors/test_tiktok_tikwm_extractor.py @@ -1,117 +1,115 @@ from datetime import datetime, timezone import time import pytest +import yt_dlp -from auto_archiver.modules.tiktok_tikwm_extractor.tiktok_tikwm_extractor import TiktokTikwmExtractor +from auto_archiver.modules.generic_extractor.generic_extractor import GenericExtractor from .test_extractor_base import TestExtractorBase +@pytest.fixture(autouse=True) +def skip_ytdlp_own_methods(mocker): + # mock this method, so that we skip the ytdlp download in these tests + mocker.patch("auto_archiver.modules.generic_extractor.tiktok.Tiktok.skip_ytdlp_download", return_value=True) + mocker.patch( + "auto_archiver.modules.generic_extractor.generic_extractor.GenericExtractor.suitable_extractors", + return_value=[e for e in yt_dlp.YoutubeDL()._ies.values() if e.IE_NAME == "TikTok"], + ) + + +@pytest.fixture() +def mock_get(mocker): + return mocker.patch("auto_archiver.modules.generic_extractor.tiktok.requests.get") + + class TestTiktokTikwmExtractor(TestExtractorBase): """ Test suite for TestTiktokTikwmExtractor. """ - extractor_module = "tiktok_tikwm_extractor" - extractor: TiktokTikwmExtractor + extractor_module = "generic_extractor" + extractor: GenericExtractor config = {} VALID_EXAMPLE_URL = "https://www.tiktok.com/@example/video/1234" - @staticmethod - def get_mockers(mocker): - mock_get = mocker.patch("auto_archiver.modules.tiktok_tikwm_extractor.tiktok_tikwm_extractor.requests.get") - mock_logger = mocker.patch("auto_archiver.modules.tiktok_tikwm_extractor.tiktok_tikwm_extractor.logger") - return mock_get, mock_logger - - @pytest.mark.parametrize("url,valid_url", [ - ("https://bellingcat.com", False), - ("https://youtube.com", False), - ("https://tiktok.co/", False), - ("https://tiktok.com/", False), - ("https://www.tiktok.com/", False), - ("https://api.cool.tiktok.com/", False), - (VALID_EXAMPLE_URL, True), - ("https://www.tiktok.com/@bbcnews/video/7478038212070411542", True), - ("https://www.tiktok.com/@ggs68taiwan.official/video/7441821351142362375", True), - ]) - def test_valid_urls(self, mocker, make_item, url, valid_url): - mock_get, mock_logger = self.get_mockers(mocker) - if valid_url: - mock_get.return_value.status_code = 404 - assert self.extractor.download(make_item(url)) == False - assert mock_get.call_count == int(valid_url) - assert mock_logger.error.call_count == int(valid_url) - - def test_invalid_json_responses(self, mocker, make_item): - mock_get, mock_logger = self.get_mockers(mocker) + def test_invalid_json_responses(self, mock_get, make_item, caplog): mock_get.return_value.status_code = 200 mock_get.return_value.json.side_effect = ValueError - assert self.extractor.download(make_item(self.VALID_EXAMPLE_URL)) == False - mock_get.assert_called_once() - mock_get.return_value.json.assert_called_once() - mock_logger.error.assert_called_once() - assert mock_logger.error.call_args[0][0].startswith("failed to parse JSON response") + with caplog.at_level("DEBUG"): + assert self.extractor.download(make_item(self.VALID_EXAMPLE_URL)) is False + mock_get.assert_called_once() + mock_get.return_value.json.assert_called_once() + # first message is just the 'Skipping using ytdlp to download files for TikTok' message + assert ( + "failed to parse JSON response from tikwm.com for url='https://www.tiktok.com/@example/video/1234'" + in caplog.text + ) mock_get.return_value.json.side_effect = Exception - with pytest.raises(Exception): - self.extractor.download(make_item(self.VALID_EXAMPLE_URL)) - mock_get.assert_called() - assert mock_get.call_count == 2 - assert mock_get.return_value.json.call_count == 2 + with caplog.at_level("ERROR"): + assert self.extractor.download(make_item(self.VALID_EXAMPLE_URL)) is False + mock_get.assert_called() + assert mock_get.call_count == 2 + assert mock_get.return_value.json.call_count == 2 + assert ( + "failed to parse JSON response from tikwm.com for url='https://www.tiktok.com/@example/video/1234'" + in caplog.text + ) - @pytest.mark.parametrize("response", [ - ({"msg": "failure"}), - ({"msg": "success"}), - ]) - def test_unsuccessful_responses(self, mocker, make_item, response): - mock_get, mock_logger = self.get_mockers(mocker) + @pytest.mark.parametrize( + "response", + [ + ({"msg": "failure"}), + ({"msg": "success"}), + ], + ) + def test_unsuccessful_responses(self, mock_get, make_item, response, caplog): mock_get.return_value.status_code = 200 mock_get.return_value.json.return_value = response - assert self.extractor.download(make_item(self.VALID_EXAMPLE_URL)) == False - mock_get.assert_called_once() - mock_get.return_value.json.assert_called_once() - mock_logger.error.assert_called_once() - assert mock_logger.error.call_args[0][0].startswith("failed to get a valid response") + with caplog.at_level("DEBUG"): + assert self.extractor.download(make_item(self.VALID_EXAMPLE_URL)) is False + mock_get.assert_called_once() + mock_get.return_value.json.assert_called_once() + assert "failed to get a valid response from tikwm.com" in caplog.text - @pytest.mark.parametrize("response,has_vid", [ - ({"data": {"id": 123}}, False), - ({"data": {"wmplay": "url"}}, True), - ({"data": {"play": "url"}}, True), - ]) - def test_correct_extraction(self, mocker, make_item, response, has_vid): - mock_get, mock_logger = self.get_mockers(mocker) + @pytest.mark.parametrize( + "response,has_vid", + [ + ({"data": {"id": 123}}, False), + ({"data": {"wmplay": "url"}}, True), + ({"data": {"play": "url"}}, True), + ], + ) + def test_correct_extraction(self, mock_get, make_item, response, has_vid, mocker): mock_get.return_value.status_code = 200 mock_get.return_value.json.return_value = {"msg": "success", **response} - result = self.extractor.download(make_item(self.VALID_EXAMPLE_URL)) if not has_vid: - assert result == False + assert result is False else: assert result.is_success() assert len(result.media) == 1 mock_get.assert_called() assert mock_get.call_count == 1 + int(has_vid) mock_get.return_value.json.assert_called_once() - if not has_vid: - mock_logger.error.assert_called_once() - assert mock_logger.error.call_args[0][0].startswith("no valid video URL found") - else: - mock_logger.error.assert_not_called() - def test_correct_extraction(self, mocker, make_item): - mock_get, _ = self.get_mockers(mocker) + def test_correct_data_extracted(self, mock_get, make_item): mock_get.return_value.status_code = 200 - mock_get.return_value.json.return_value = {"msg": "success", "data": { - "wmplay": "url", - "origin_cover": "cover.jpg", - "title": "Title", - "id": 123, - "duration": 60, - "create_time": 1736301699, - "author": "Author", - "other": "data" - }} + mock_get.return_value.json.return_value = { + "msg": "success", + "data": { + "wmplay": "url", + "origin_cover": "cover.jpg", + "title": "Title", + "id": 123, + "duration": 60, + "create_time": 1736301699, + "author": "Author", + "other": "data", + }, + } result = self.extractor.download(make_item(self.VALID_EXAMPLE_URL)) assert result.is_success() @@ -129,26 +127,25 @@ class TestTiktokTikwmExtractor(TestExtractorBase): result = self.extractor.download(make_item(url)) assert result.is_success() assert len(result.media) == 2 - assert result.get_title() == "The A23a iceberg is one of the world's oldest and it's so big you can see it from space. #Iceberg #A23a #Antarctica #Ice #ClimateChange #DavidAttenborough #Ocean #Sea #SouthGeorgia #BBCNews " + assert ( + result.get_title() + == "The A23a iceberg is one of the world's oldest and it's so big you can see it from space. #Iceberg #A23a #Antarctica #Ice #ClimateChange #DavidAttenborough #Ocean #Sea #SouthGeorgia #BBCNews " + ) assert result.get("author").get("unique_id") == "bbcnews" - assert result.get("api_data").get("id") == '7478038212070411542' + assert result.get("api_data").get("id") == "7478038212070411542" assert result.media[1].get("duration") == 59 assert result.get("timestamp") == datetime.fromtimestamp(1741122000, tz=timezone.utc) @pytest.mark.download - def test_download_sensitive_video(self, make_item, mock_sleep): - # sleep is needed because of the rate limit - mock_sleep.stop() - time.sleep(1.1) - mock_sleep.start() - + def test_download_sensitive_video(self, make_item): url = "https://www.tiktok.com/@ggs68taiwan.official/video/7441821351142362375" - + # Required for rate limiting + time.sleep(1.1) result = self.extractor.download(make_item(url)) assert result.is_success() assert len(result.media) == 2 assert result.get_title() == "Căng nhất lúc này #ggs68 #ggs68taiwan #taiwan #dailoan #tiktoknews" assert result.get("author").get("id") == "7197400619475649562" - assert result.get("api_data").get("id") == '7441821351142362375' + assert result.get("api_data").get("id") == "7441821351142362375" assert result.media[1].get("duration") == 34 assert result.get("timestamp") == datetime.fromtimestamp(1732684060, tz=timezone.utc) diff --git a/tests/extractors/test_twitter_api_extractor.py b/tests/extractors/test_twitter_api_extractor.py index 26394ac..8b8e0d9 100644 --- a/tests/extractors/test_twitter_api_extractor.py +++ b/tests/extractors/test_twitter_api_extractor.py @@ -1,6 +1,5 @@ import os import datetime -import hashlib import pytest from pytwitter.models.media import MediaVariant @@ -10,8 +9,7 @@ from auto_archiver.modules.twitter_api_extractor import TwitterApiExtractor @pytest.mark.incremental class TestTwitterApiExtractor(TestExtractorBase): - - extractor_module = 'twitter_api_extractor' + extractor_module: TwitterApiExtractor = "twitter_api_extractor" config = { "bearer_tokens": [], @@ -22,41 +20,79 @@ class TestTwitterApiExtractor(TestExtractorBase): "access_secret": os.environ.get("TWITTER_ACCESS_SECRET"), } - @pytest.mark.parametrize("url, expected", [ - ("https://x.com/bellingcat/status/1874097816571961839", "https://x.com/bellingcat/status/1874097816571961839"), # x.com urls unchanged - ("https://twitter.com/bellingcat/status/1874097816571961839", "https://twitter.com/bellingcat/status/1874097816571961839"), # twitter urls unchanged - ("https://twitter.com/bellingcat/status/1874097816571961839?s=20&t=3d0g4ZQis7dCbSDg-mE7-w", "https://twitter.com/bellingcat/status/1874097816571961839?s=20&t=3d0g4ZQis7dCbSDg-mE7-w"), # don't strip params from twitter urls (changed Jan 2025) - ("https://www.bellingcat.com/category/resources/", "https://www.bellingcat.com/category/resources/"), # non-twitter/x urls unchanged - ("https://www.bellingcat.com/category/resources/?s=20&t=3d0g4ZQis7dCbSDg-mE7-w", "https://www.bellingcat.com/category/resources/?s=20&t=3d0g4ZQis7dCbSDg-mE7-w"), # shouldn't strip params from non-twitter/x URLs - ]) + @pytest.mark.parametrize( + "url, expected", + [ + ( + "https://x.com/bellingcat/status/1874097816571961839", + "https://x.com/bellingcat/status/1874097816571961839", + ), # x.com urls unchanged + ( + "https://twitter.com/bellingcat/status/1874097816571961839", + "https://twitter.com/bellingcat/status/1874097816571961839", + ), # twitter urls unchanged + ( + "https://twitter.com/bellingcat/status/1874097816571961839?s=20&t=3d0g4ZQis7dCbSDg-mE7-w", + "https://twitter.com/bellingcat/status/1874097816571961839?s=20&t=3d0g4ZQis7dCbSDg-mE7-w", + ), # don't strip params from twitter urls (changed Jan 2025) + ( + "https://www.bellingcat.com/category/resources/", + "https://www.bellingcat.com/category/resources/", + ), # non-twitter/x urls unchanged + ( + "https://www.bellingcat.com/category/resources/?s=20&t=3d0g4ZQis7dCbSDg-mE7-w", + "https://www.bellingcat.com/category/resources/?s=20&t=3d0g4ZQis7dCbSDg-mE7-w", + ), # shouldn't strip params from non-twitter/x URLs + ], + ) def test_sanitize_url(self, url, expected): assert expected == self.extractor.sanitize_url(url) @pytest.mark.download def test_sanitize_url_download(self): - assert "https://www.bellingcat.com/category/resources/" == self.extractor.sanitize_url("https://t.co/yl3oOJatFp") + assert "https://www.bellingcat.com/category/resources/" == self.extractor.sanitize_url( + "https://t.co/yl3oOJatFp" + ) - @pytest.mark.parametrize("url, exptected_username, exptected_tweetid", [ - ("https://twitter.com/bellingcat/status/1874097816571961839", "bellingcat", "1874097816571961839"), - ("https://x.com/bellingcat/status/1874097816571961839", "bellingcat", "1874097816571961839"), - ("https://www.bellingcat.com/category/resources/", False, False) - ]) + @pytest.mark.parametrize( + "url, exptected_username, exptected_tweetid", + [ + ("https://twitter.com/bellingcat/status/1874097816571961839", "bellingcat", "1874097816571961839"), + ("https://x.com/bellingcat/status/1874097816571961839", "bellingcat", "1874097816571961839"), + ("https://www.bellingcat.com/category/resources/", False, False), + ], + ) def test_get_username_tweet_id_from_url(self, url, exptected_username, exptected_tweetid): - username, tweet_id = self.extractor.get_username_tweet_id(url) assert exptected_username == username assert exptected_tweetid == tweet_id def test_choose_variants(self): # taken from the response for url https://x.com/bellingcat/status/1871552600346415571 - variant_list = [MediaVariant(content_type='application/x-mpegURL', url='https://video.twimg.com/ext_tw_video/1871551993677852672/pu/pl/ovWo7ux-bKROwYIC.m3u8?tag=12&v=e1b'), - MediaVariant(bit_rate=256000, content_type='video/mp4', url='https://video.twimg.com/ext_tw_video/1871551993677852672/pu/vid/avc1/480x270/OqZIrKV0LFswMvxS.mp4?tag=12'), - MediaVariant(bit_rate=832000, content_type='video/mp4', url='https://video.twimg.com/ext_tw_video/1871551993677852672/pu/vid/avc1/640x360/uiDZDSmZ8MZn9hsi.mp4?tag=12'), - MediaVariant(bit_rate=2176000, content_type='video/mp4', url='https://video.twimg.com/ext_tw_video/1871551993677852672/pu/vid/avc1/1280x720/6Y340Esh568WZnRZ.mp4?tag=12') - ] + variant_list = [ + MediaVariant( + content_type="application/x-mpegURL", + url="https://video.twimg.com/ext_tw_video/1871551993677852672/pu/pl/ovWo7ux-bKROwYIC.m3u8?tag=12&v=e1b", + ), + MediaVariant( + bit_rate=256000, + content_type="video/mp4", + url="https://video.twimg.com/ext_tw_video/1871551993677852672/pu/vid/avc1/480x270/OqZIrKV0LFswMvxS.mp4?tag=12", + ), + MediaVariant( + bit_rate=832000, + content_type="video/mp4", + url="https://video.twimg.com/ext_tw_video/1871551993677852672/pu/vid/avc1/640x360/uiDZDSmZ8MZn9hsi.mp4?tag=12", + ), + MediaVariant( + bit_rate=2176000, + content_type="video/mp4", + url="https://video.twimg.com/ext_tw_video/1871551993677852672/pu/vid/avc1/1280x720/6Y340Esh568WZnRZ.mp4?tag=12", + ), + ] chosen_variant = self.extractor.choose_variant(variant_list) assert chosen_variant == variant_list[3] - + @pytest.mark.skipif(not os.environ.get("TWITTER_BEARER_TOKEN"), reason="No Twitter bearer token provided") @pytest.mark.download def test_download_nonexistent_tweet(self, make_item): @@ -76,7 +112,6 @@ class TestTwitterApiExtractor(TestExtractorBase): @pytest.mark.skipif(not os.environ.get("TWITTER_BEARER_TOKEN"), reason="No Twitter bearer token provided") @pytest.mark.download def test_download_tweet_no_media(self, make_item): - item = make_item("https://twitter.com/MeCookieMonster/status/1617921633456640001?s=20&t=3d0g4ZQis7dCbSDg-mE7-w") post = self.extractor.download(item) @@ -84,7 +119,7 @@ class TestTwitterApiExtractor(TestExtractorBase): post, "Onion rings are just vegetable donuts.", datetime.datetime(2023, 1, 24, 16, 25, 51, tzinfo=datetime.timezone.utc), - "twitter-api: success" + "twitter-api: success", ) @pytest.mark.skipif(not os.environ.get("TWITTER_BEARER_TOKEN"), reason="No Twitter bearer token provided") @@ -95,27 +130,41 @@ class TestTwitterApiExtractor(TestExtractorBase): self.assertValidResponseMetadata( post, "This month's Bellingchat Premium is with @KolinaKoltai. She reveals how she investigated a platform allowing users to create AI-generated child sexual abuse material and explains why it's crucial to investigate the people behind these services https://t.co/SfBUq0hSD0 https://t.co/rIHx0WlKp8", - datetime.datetime(2024, 12, 24, 13, 44, 46, tzinfo=datetime.timezone.utc) + datetime.datetime(2024, 12, 24, 13, 44, 46, tzinfo=datetime.timezone.utc), ) @pytest.mark.skipif(not os.environ.get("TWITTER_BEARER_TOKEN"), reason="No Twitter bearer token provided") - @pytest.mark.parametrize("url, title, timestamp", [ - ("https://x.com/SozinhoRamalho/status/1876710769913450647", "ignore tweet, testing sensitivity warning nudity https://t.co/t3u0hQsSB1", datetime.datetime(2024, 12, 31, 14, 18, 33, tzinfo=datetime.timezone.utc)), - ("https://x.com/SozinhoRamalho/status/1876710875475681357", "ignore tweet, testing sensitivity warning violence https://t.co/syYDSkpjZD", datetime.datetime(2024, 12, 31, 14, 18, 33, tzinfo=datetime.timezone.utc)), - ("https://x.com/SozinhoRamalho/status/1876711053813227618", "ignore tweet, testing sensitivity warning sensitive https://t.co/XE7cRdjzYq", datetime.datetime(2024, 12, 31, 14, 18, 33, tzinfo=datetime.timezone.utc)), - ("https://x.com/SozinhoRamalho/status/1876711141314801937", "ignore tweet, testing sensitivity warning nudity, violence, sensitivity https://t.co/YxCFbbhYE3", datetime.datetime(2024, 12, 31, 14, 18, 33, tzinfo=datetime.timezone.utc)), - ]) + @pytest.mark.parametrize( + "url, title, timestamp", + [ + ( + "https://x.com/SozinhoRamalho/status/1876710769913450647", + "ignore tweet, testing sensitivity warning nudity https://t.co/t3u0hQsSB1", + datetime.datetime(2024, 12, 31, 14, 18, 33, tzinfo=datetime.timezone.utc), + ), + ( + "https://x.com/SozinhoRamalho/status/1876710875475681357", + "ignore tweet, testing sensitivity warning violence https://t.co/syYDSkpjZD", + datetime.datetime(2024, 12, 31, 14, 18, 33, tzinfo=datetime.timezone.utc), + ), + ( + "https://x.com/SozinhoRamalho/status/1876711053813227618", + "ignore tweet, testing sensitivity warning sensitive https://t.co/XE7cRdjzYq", + datetime.datetime(2024, 12, 31, 14, 18, 33, tzinfo=datetime.timezone.utc), + ), + ( + "https://x.com/SozinhoRamalho/status/1876711141314801937", + "ignore tweet, testing sensitivity warning nudity, violence, sensitivity https://t.co/YxCFbbhYE3", + datetime.datetime(2024, 12, 31, 14, 18, 33, tzinfo=datetime.timezone.utc), + ), + ], + ) @pytest.mark.download def test_download_sensitive_media(self, url, title, timestamp, check_hash, make_item): - """Download tweets with sensitive media""" post = self.extractor.download(make_item(url)) - self.assertValidResponseMetadata( - post, - title, - timestamp - ) + self.assertValidResponseMetadata(post, title, timestamp) assert len(post.media) == 1 # check the SHA1 hash (quick) of the media, to make sure it's valid - check_hash(post.media[0].filename, "3eea9c03b2dcedd1eb9a169d8bfd1cf877996fab4961de019a96eb9d32d2d733") \ No newline at end of file + check_hash(post.media[0].filename, "3eea9c03b2dcedd1eb9a169d8bfd1cf877996fab4961de019a96eb9d32d2d733") diff --git a/tests/extractors/test_vk_extractor.py b/tests/extractors/test_vk_extractor.py index 80eb9dd..040e5f7 100644 --- a/tests/extractors/test_vk_extractor.py +++ b/tests/extractors/test_vk_extractor.py @@ -9,6 +9,7 @@ def mock_vk_scraper(mocker): """Fixture to mock VkScraper.""" return mocker.patch("auto_archiver.modules.vk_extractor.vk_extractor.VkScraper") + @pytest.fixture def vk_extractor(setup_module, mock_vk_scraper) -> VkExtractor: """Fixture to initialize VkExtractor with mocked VkScraper.""" @@ -39,7 +40,7 @@ def test_vk_url_but_scrape_returns_empty(vk_extractor, metadata): def test_successful_scrape_and_download(vk_extractor, metadata, mocker): mock_scrapes = [ {"text": "Post Title", "datetime": "2023-01-01T00:00:00", "id": 1}, - {"text": "Another Post", "datetime": "2023-01-02T00:00:00", "id": 2} + {"text": "Another Post", "datetime": "2023-01-02T00:00:00", "id": 2}, ] mock_filenames = ["image1.jpg", "image2.png"] vk_extractor.vks.scrape.return_value = mock_scrapes @@ -56,16 +57,16 @@ def test_successful_scrape_and_download(vk_extractor, metadata, mocker): assert len(result.media) == 2 assert result.media[0].filename == "image1.jpg" assert result.media[1].filename == "image2.png" - vk_extractor.vks.download_media.assert_called_once_with( - mock_scrapes, vk_extractor.tmp_dir - ) + vk_extractor.vks.download_media.assert_called_once_with(mock_scrapes, vk_extractor.tmp_dir) def test_adds_first_title_and_timestamp(vk_extractor): metadata = Metadata().set_url("https://vk.com/no-metadata") metadata.set_url("https://vk.com/no-metadata") - mock_scrapes = [{"text": "value", "datetime": "2023-01-01T00:00:00"}, - {"text": "value2", "datetime": "2023-01-02T00:00:00"}] + mock_scrapes = [ + {"text": "value", "datetime": "2023-01-01T00:00:00"}, + {"text": "value2", "datetime": "2023-01-02T00:00:00"}, + ] vk_extractor.vks.scrape.return_value = mock_scrapes vk_extractor.vks.download_media.return_value = [] result = vk_extractor.download(metadata) @@ -73,4 +74,4 @@ def test_adds_first_title_and_timestamp(vk_extractor): assert result.get_title() == "value" # formatted timestamp assert result.get_timestamp() == "2023-01-01T00:00:00+00:00" - assert result.is_success() \ No newline at end of file + assert result.is_success() diff --git a/tests/feeders/test_atlos_feeder.py b/tests/feeders/test_atlos_feeder.py index 1ef9fab..f423136 100644 --- a/tests/feeders/test_atlos_feeder.py +++ b/tests/feeders/test_atlos_feeder.py @@ -36,29 +36,45 @@ def atlos_feeder(setup_module, mocker) -> AtlosFeeder: @pytest.fixture def mock_atlos_api(atlos_feeder): """Fixture to update the atlos_feeder.session.get side_effect.""" + def _mock_responses(responses): atlos_feeder.session.get.side_effect = [FakeAPIResponse(data) for data in responses] + return _mock_responses def test_atlos_feeder_iter_yields_valid_metadata(atlos_feeder, mock_atlos_api): """Test valid items are yielded and invalid ones ignored.""" - mock_atlos_api([ - { - "next": None, - "results": [ - {"source_url": "http://example.com", "id": 1, - "metadata": {"auto_archiver": {"processed": False}}, - "visibility": "visible", "status": "complete"}, - {"source_url": "", "id": 2, - "metadata": {"auto_archiver": {"processed": False}}, - "visibility": "visible", "status": "complete"}, - {"source_url": "http://example.org", "id": 3, - "metadata": {"auto_archiver": {"processed": True}}, - "visibility": "visible", "status": "complete"}, - ], - } - ]) + mock_atlos_api( + [ + { + "next": None, + "results": [ + { + "source_url": "http://example.com", + "id": 1, + "metadata": {"auto_archiver": {"processed": False}}, + "visibility": "visible", + "status": "complete", + }, + { + "source_url": "", + "id": 2, + "metadata": {"auto_archiver": {"processed": False}}, + "visibility": "visible", + "status": "complete", + }, + { + "source_url": "http://example.org", + "id": 3, + "metadata": {"auto_archiver": {"processed": True}}, + "visibility": "visible", + "status": "complete", + }, + ], + } + ] + ) items = list(atlos_feeder) assert len(items) == 1 @@ -68,24 +84,34 @@ def test_atlos_feeder_iter_yields_valid_metadata(atlos_feeder, mock_atlos_api): def test_atlos_feeder_multiple_pages(atlos_feeder, mock_atlos_api): """Test iteration over multiple pages with valid items.""" - mock_atlos_api([ - { - "next": "cursor2", - "results": [ - {"source_url": "http://example1.com", "id": 10, - "metadata": {"auto_archiver": {"processed": False}}, - "visibility": "visible", "status": "complete"}, - ], - }, - { - "next": None, - "results": [ - {"source_url": "http://example2.com", "id": 20, - "metadata": {"auto_archiver": {"processed": False}}, - "visibility": "visible", "status": "complete"}, - ], - }, - ]) + mock_atlos_api( + [ + { + "next": "cursor2", + "results": [ + { + "source_url": "http://example1.com", + "id": 10, + "metadata": {"auto_archiver": {"processed": False}}, + "visibility": "visible", + "status": "complete", + }, + ], + }, + { + "next": None, + "results": [ + { + "source_url": "http://example2.com", + "id": 20, + "metadata": {"auto_archiver": {"processed": False}}, + "visibility": "visible", + "status": "complete", + }, + ], + }, + ] + ) items = list(atlos_feeder) assert len(items) == 2 diff --git a/tests/feeders/test_csv_feeder.py b/tests/feeders/test_csv_feeder.py index 546c3a7..965f8ad 100644 --- a/tests/feeders/test_csv_feeder.py +++ b/tests/feeders/test_csv_feeder.py @@ -1,13 +1,16 @@ import pytest + @pytest.fixture def headerless_csv_file(): return "tests/data/csv_no_headers.csv" + @pytest.fixture def header_csv_file(): return "tests/data/csv_with_headers.csv" + @pytest.fixture def header_csv_file_non_default_column(): return "tests/data/csv_with_headers_non_default_column.csv" @@ -23,6 +26,7 @@ def test_csv_feeder_no_headers(headerless_csv_file, setup_module): assert urls[0].get_url() == "https://example.com/1/" assert urls[1].get_url() == "https://example.com/2/" + def test_csv_feeder_with_headers(header_csv_file, setup_module): from auto_archiver.modules.csv_feeder.csv_feeder import CSVFeeder @@ -33,10 +37,10 @@ def test_csv_feeder_with_headers(header_csv_file, setup_module): assert urls[0].get_url() == "https://example.com/1/" assert urls[1].get_url() == "https://example.com/2/" + def test_csv_feeder_wrong_column(header_csv_file, setup_module, caplog): from auto_archiver.modules.csv_feeder.csv_feeder import CSVFeeder - with caplog.at_level("WARNING"): feeder = setup_module(CSVFeeder, {"files": [header_csv_file], "column": 1}) urls = list(feeder) @@ -54,4 +58,4 @@ def test_csv_feeder_column_by_name(header_csv_file, setup_module): urls = list(feeder) assert len(urls) == 2 assert urls[0].get_url() == "https://example.com/1/" - assert urls[1].get_url() == "https://example.com/2/" \ No newline at end of file + assert urls[1].get_url() == "https://example.com/2/" diff --git a/tests/feeders/test_gsheet_feeder.py b/tests/feeders/test_gsheet_feeder.py index 9ca81b0..bf34757 100644 --- a/tests/feeders/test_gsheet_feeder.py +++ b/tests/feeders/test_gsheet_feeder.py @@ -19,35 +19,32 @@ def test_setup_without_sheet_and_sheet_id(setup_module, mocker): @pytest.fixture def gsheet_feeder(setup_module, mocker) -> GsheetsFeederDB: config: dict = { - "service_account": "dummy.json", - "sheet": "test-auto-archiver", - "sheet_id": None, - "header": 1, - "columns": { - "url": "link", - "status": "archive status", - "folder": "destination folder", - "archive": "archive location", - "date": "archive date", - "thumbnail": "thumbnail", - "timestamp": "upload timestamp", - "title": "upload title", - "text": "text content", - "screenshot": "screenshot", - "hash": "hash", - "pdq_hash": "perceptual hashes", - "wacz": "wacz", - "replaywebpage": "replaywebpage", - }, - "allow_worksheets": set(), - "block_worksheets": set(), - "use_sheet_names_in_stored_paths": True, - } + "service_account": "dummy.json", + "sheet": "test-auto-archiver", + "sheet_id": None, + "header": 1, + "columns": { + "url": "link", + "status": "archive status", + "folder": "destination folder", + "archive": "archive location", + "date": "archive date", + "thumbnail": "thumbnail", + "timestamp": "upload timestamp", + "title": "upload title", + "text": "text content", + "screenshot": "screenshot", + "hash": "hash", + "pdq_hash": "perceptual hashes", + "wacz": "wacz", + "replaywebpage": "replaywebpage", + }, + "allow_worksheets": set(), + "block_worksheets": set(), + "use_sheet_names_in_stored_paths": True, + } mocker.patch("gspread.service_account") - feeder = setup_module( - "gsheet_feeder_db", - config - ) + feeder = setup_module("gsheet_feeder_db", config) feeder.gsheets_client = mocker.MagicMock() return feeder @@ -128,9 +125,7 @@ def test__set_metadata_with_folder(gsheet_feeder: GsheetsFeederDB): (None, "ABC123", "open_by_key", "ABC123", "opening by sheet ID"), ], ) -def test_open_sheet_with_name_or_id( - setup_module, sheet, sheet_id, expected_method, expected_arg, description, mocker -): +def test_open_sheet_with_name_or_id(setup_module, sheet, sheet_id, expected_method, expected_arg, description, mocker): """Ensure open_sheet() correctly opens by name or ID based on configuration.""" mock_service_account = mocker.patch("gspread.service_account") mock_client = mocker.MagicMock() @@ -145,9 +140,7 @@ def test_open_sheet_with_name_or_id( ) sheet_result = feeder.open_sheet() # Validate the correct method was called - getattr(mock_client, expected_method).assert_called_once_with( - expected_arg - ), f"Failed: {description}" + getattr(mock_client, expected_method).assert_called_once_with(expected_arg), f"Failed: {description}" assert sheet_result == "MockSheet", f"Failed: {description}" @@ -179,10 +172,10 @@ def test_should_process_sheet(setup_module, mocker): "block_worksheets": {"Sheet3"}, }, ) - assert gdb.should_process_sheet("TestSheet") == True - assert gdb.should_process_sheet("Sheet3") == False + assert gdb.should_process_sheet("TestSheet") is True + assert gdb.should_process_sheet("Sheet3") is False # False if allow_worksheets is set - assert gdb.should_process_sheet("AnotherSheet") == False + assert gdb.should_process_sheet("AnotherSheet") is False @pytest.mark.skip(reason="Requires a real connection") @@ -220,9 +213,7 @@ class TestGSheetsFeederReal: @pytest.fixture(autouse=True) def setup_feeder(self, setup_module): - assert ( - self.module_name is not None - ), "self.module_name must be set on the subclass" + assert self.module_name is not None, "self.module_name must be set on the subclass" assert self.config is not None, "self.config must be a dict set on the subclass" self.feeder: Type[Feeder] = setup_module(self.module_name, self.config) @@ -241,9 +232,7 @@ class TestGSheetsFeederReal: """Ensure open_sheet() connects to a real Google Sheets instance.""" sheet = self.feeder.open_sheet() assert sheet is not None, "open_sheet() should return a valid sheet instance" - assert hasattr( - sheet, "worksheets" - ), "Returned object should have worksheets method" + assert hasattr(sheet, "worksheets"), "Returned object should have worksheets method" def test_iter_yields_metadata_real_data(self): """Ensure __iter__() yields Metadata objects for real test sheet data.""" diff --git a/tests/feeders/test_gworksheet.py b/tests/feeders/test_gworksheet.py index b6a0b5c..c817be3 100644 --- a/tests/feeders/test_gworksheet.py +++ b/tests/feeders/test_gworksheet.py @@ -81,40 +81,27 @@ class TestGWorksheet: (False, ""), ], ) - def test_get_cell_or_default_handles_empty_values( - self, mock_worksheet, when_empty, expected - ): + def test_get_cell_or_default_handles_empty_values(self, mock_worksheet, when_empty, expected): mock_worksheet.get_values.return_value[1][0] = "" # Empty URL cell g = GWorksheet(mock_worksheet) - assert ( - g.get_cell_or_default( - 2, "url", default="default", when_empty_use_default=when_empty - ) - == expected - ) + assert g.get_cell_or_default(2, "url", default="default", when_empty_use_default=when_empty) == expected def test_get_cell_or_default_handles_missing_columns(self, gworksheet): - assert ( - gworksheet.get_cell_or_default(1, "invalid_col", default="safe") == "safe" - ) + assert gworksheet.get_cell_or_default(1, "invalid_col", default="safe") == "safe" # Test write operations def test_set_cell_updates_correct_position(self, mock_worksheet, gworksheet): gworksheet.set_cell(2, "url", "new_url") mock_worksheet.update_cell.assert_called_once_with(2, 1, "new_url") - def test_batch_set_cell_formats_requests_correctly( - self, mock_worksheet, gworksheet - ): + def test_batch_set_cell_formats_requests_correctly(self, mock_worksheet, gworksheet): updates = [(2, "url", "new_url"), (3, "status", "processed")] gworksheet.batch_set_cell(updates) expected_batch = [ {"range": "A2", "values": [["new_url"]]}, {"range": "B3", "values": [["processed"]]}, ] - mock_worksheet.batch_update.assert_called_once_with( - expected_batch, value_input_option="USER_ENTERED" - ) + mock_worksheet.batch_update.assert_called_once_with(expected_batch, value_input_option="USER_ENTERED") def test_batch_set_cell_truncates_long_values(self, mock_worksheet, gworksheet): long_value = "x" * 50000 diff --git a/tests/formatters/test_html_formatter.py b/tests/formatters/test_html_formatter.py index 60abaa7..502e231 100644 --- a/tests/formatters/test_html_formatter.py +++ b/tests/formatters/test_html_formatter.py @@ -5,13 +5,13 @@ from auto_archiver.core import Metadata, Media def test_format(setup_module): formatter = setup_module(HtmlFormatter) - metadata = Metadata().set("content", "Hello, world!").set_url('https://example.com') + metadata = Metadata().set("content", "Hello, world!").set_url("https://example.com") final_media = formatter.format(metadata) assert isinstance(final_media, Media) assert ".html" in final_media.filename - with open (final_media.filename, "r", encoding="utf-8") as f: + with open(final_media.filename, "r", encoding="utf-8") as f: content = f.read() assert "Hello, world!" in content assert final_media.mimetype == "text/html" - assert "SHA-256:" in final_media.get('hash') \ No newline at end of file + assert "SHA-256:" in final_media.get("hash") diff --git a/tests/storages/test_S3_storage.py b/tests/storages/test_S3_storage.py index abf9763..87da776 100644 --- a/tests/storages/test_S3_storage.py +++ b/tests/storages/test_S3_storage.py @@ -8,6 +8,7 @@ class TestS3Storage: """ Test suite for S3Storage. """ + module_name: str = "s3_storage" storage: Type[S3Storage] config: dict = { @@ -32,10 +33,10 @@ class TestS3Storage: """Test that S3 client is initialized with correct parameters""" assert self.storage.s3 is not None - assert self.storage.s3.meta.region_name == 'test-region' + assert self.storage.s3.meta.region_name == "test-region" def test_get_cdn_url_generation(self): - """Test CDN URL formatting """ + """Test CDN URL formatting""" media = Media("test.txt") media._key = "path/to/file.txt" url = self.storage.get_cdn_url(media) @@ -46,14 +47,14 @@ class TestS3Storage: def test_uploadf_sets_acl_public(self, mocker): media = Media("test.txt") mock_file = mocker.MagicMock() - mock_s3_upload = mocker.patch.object(self.storage.s3, 'upload_fileobj') - mocker.patch.object(self.storage, 'is_upload_needed', return_value=True) + mock_s3_upload = mocker.patch.object(self.storage.s3, "upload_fileobj") + mocker.patch.object(self.storage, "is_upload_needed", return_value=True) self.storage.uploadf(mock_file, media) mock_s3_upload.assert_called_once_with( mock_file, - Bucket='test-bucket', + Bucket="test-bucket", Key=media.key, - ExtraArgs={'ACL': 'public-read', 'ContentType': 'text/plain'} + ExtraArgs={"ACL": "public-read", "ContentType": "text/plain"}, ) def test_upload_decision_logic(self, mocker): @@ -61,23 +62,29 @@ class TestS3Storage: media = Media("test.txt") assert self.storage.is_upload_needed(media) is True self.storage.random_no_duplicate = True - mock_calc_hash = mocker.patch('auto_archiver.modules.s3_storage.s3_storage.calculate_file_hash', return_value='beepboop123beepboop123beepboop123') - mock_file_in_folder = mocker.patch.object(self.storage, 'file_in_folder', return_value='existing_key.txt') + mocker.patch( + "auto_archiver.modules.s3_storage.s3_storage.calculate_file_hash", + return_value="beepboop123beepboop123beepboop123", + ) + mock_file_in_folder = mocker.patch.object(self.storage, "file_in_folder", return_value="existing_key.txt") assert self.storage.is_upload_needed(media) is False - assert media.key == 'existing_key.txt' - mock_file_in_folder.assert_called_with('no-dups/beepboop123beepboop123be') + assert media.key == "existing_key.txt" + mock_file_in_folder.assert_called_with("no-dups/beepboop123beepboop123be") def test_skips_upload_when_duplicate_exists(self, mocker): """Test that upload skips when file_in_folder finds existing object""" self.storage.random_no_duplicate = True - mock_file_in_folder = mocker.patch.object(S3Storage, 'file_in_folder', return_value="existing_folder/existing_file.txt") + mocker.patch.object(S3Storage, "file_in_folder", return_value="existing_folder/existing_file.txt") media = Media("test.txt") media._key = "original_path.txt" - mock_calculate_hash = mocker.patch('auto_archiver.modules.s3_storage.s3_storage.calculate_file_hash', return_value="beepboop123beepboop123beepboop123") + mocker.patch( + "auto_archiver.modules.s3_storage.s3_storage.calculate_file_hash", + return_value="beepboop123beepboop123beepboop123", + ) assert self.storage.is_upload_needed(media) is False assert media.key == "existing_folder/existing_file.txt" assert media.get("previously archived") is True - mock_upload = mocker.patch.object(self.storage.s3, 'upload_fileobj') + mock_upload = mocker.patch.object(self.storage.s3, "upload_fileobj") result = self.storage.uploadf(None, media) mock_upload.assert_not_called() assert result is True @@ -85,21 +92,18 @@ class TestS3Storage: def test_uploads_with_correct_parameters(self, mocker): media = Media("test.txt") media._key = "original_key.txt" - mocker.patch.object(S3Storage, 'is_upload_needed', return_value=True) - media.mimetype = 'image/png' + mocker.patch.object(S3Storage, "is_upload_needed", return_value=True) + media.mimetype = "image/png" mock_file = mocker.MagicMock() - mock_upload = mocker.patch.object(self.storage.s3, 'upload_fileobj') + mock_upload = mocker.patch.object(self.storage.s3, "upload_fileobj") self.storage.uploadf(mock_file, media) mock_upload.assert_called_once_with( mock_file, - Bucket='test-bucket', - Key='original_key.txt', - ExtraArgs={ - 'ACL': 'public-read', - 'ContentType': 'image/png' - } + Bucket="test-bucket", + Key="original_key.txt", + ExtraArgs={"ACL": "public-read", "ContentType": "image/png"}, ) def test_file_in_folder_exists(self, mocker): - mock_list_objects = mocker.patch.object(self.storage.s3, 'list_objects', return_value={'Contents': [{'Key': 'path/to/file.txt'}]}) - assert self.storage.file_in_folder('path/to/') == 'path/to/file.txt' + mocker.patch.object(self.storage.s3, "list_objects", return_value={"Contents": [{"Key": "path/to/file.txt"}]}) + assert self.storage.file_in_folder("path/to/") == "path/to/file.txt" diff --git a/tests/storages/test_atlos_storage.py b/tests/storages/test_atlos_storage.py index d33ff2f..f273c7e 100644 --- a/tests/storages/test_atlos_storage.py +++ b/tests/storages/test_atlos_storage.py @@ -90,18 +90,19 @@ def test_upload_not_uploaded(tmp_path, atlos_storage: AtlosStorage, metadata: Me get_mock.assert_called_once() post_mock.assert_called_once() - expected_endpoint = f"/api/v2/source_material/upload/202" + expected_endpoint = "/api/v2/source_material/upload/202" call_args = post_mock.call_args[0] assert call_args[0] == expected_endpoint call_kwargs = post_mock.call_args[1] - expected_headers = {"Authorization": f"Bearer {atlos_storage.api_token}"} expected_params = {"title": media.properties} assert call_kwargs["params"] == expected_params file_tuple = call_kwargs["files"]["file"] assert file_tuple[0] == os.path.basename(media.filename) -def test_upload_post_http_error(tmp_path, atlos_storage: AtlosStorage, metadata: Metadata, media: Media, mocker) -> None: +def test_upload_post_http_error( + tmp_path, atlos_storage: AtlosStorage, metadata: Metadata, media: Media, mocker +) -> None: """Test upload() propagates HTTP error during POST.""" metadata.set("atlos_id", 303) fake_get_response = {"result": {"artifacts": []}} @@ -109,4 +110,3 @@ def test_upload_post_http_error(tmp_path, atlos_storage: AtlosStorage, metadata: mocker.patch.object(atlos_storage, "_post", side_effect=Exception("HTTP error")) with pytest.raises(Exception, match="HTTP error"): atlos_storage.upload(media, metadata) - diff --git a/tests/storages/test_gdrive_storage.py b/tests/storages/test_gdrive_storage.py index d87f5e8..99df536 100644 --- a/tests/storages/test_gdrive_storage.py +++ b/tests/storages/test_gdrive_storage.py @@ -1,37 +1,42 @@ from typing import Type import pytest -from oauth2client import service_account from auto_archiver.core import Media from auto_archiver.modules.gdrive_storage import GDriveStorage -from auto_archiver.core.metadata import Metadata from tests.storages.test_storage_base import TestStorageBase +@pytest.fixture(autouse=True) +def mock_sleep(mocker): + """Mock time.sleep to avoid delays.""" + return mocker.patch("time.sleep") + + @pytest.fixture -def gdrive_storage(setup_module, mocker): +def gdrive_storage(setup_module, mocker) -> GDriveStorage: module_name: str = "gdrive_storage" - storage: GDriveStorage - config: dict = {'path_generator': 'url', - 'filename_generator': 'static', - 'root_folder_id': "fake_root_folder_id", - 'oauth_token': None, - 'service_account': 'fake_service_account.json' - } - mocker.patch('google.oauth2.service_account.Credentials.from_service_account_file') + config: dict = { + "path_generator": "url", + "filename_generator": "static", + "root_folder_id": "fake_root_folder_id", + "oauth_token": None, + "service_account": "fake_service_account.json", + } + mocker.patch("google.oauth2.service_account.Credentials.from_service_account_file") return setup_module(module_name, config) def test_initialize_fails_with_non_existent_creds(setup_module): """Test that the Google Drive service raises a FileNotFoundError when the service account file does not exist. - (and isn't mocked) + (and isn't mocked) """ - config: dict = {'path_generator': 'url', - 'filename_generator': 'static', - 'root_folder_id': "fake_root_folder_id", - 'oauth_token': None, - 'service_account': 'fake_service_account.json' - } + config: dict = { + "path_generator": "url", + "filename_generator": "static", + "root_folder_id": "fake_root_folder_id", + "oauth_token": None, + "service_account": "fake_service_account.json", + } with pytest.raises(FileNotFoundError) as exc_info: setup_module("gdrive_storage", config) assert "No such file or directory" in str(exc_info.value) @@ -48,12 +53,12 @@ def test_get_id_from_parent_and_name(gdrive_storage, mocker): result = gdrive_storage._get_id_from_parent_and_name("parent", "mock", retries=1, use_mime_type=False) assert result == "123" + def test_path_parts(): media = Media(filename="test.jpg") media._key = "folder1/folder2/test.jpg" - @pytest.mark.skip(reason="Requires real credentials") @pytest.mark.download class TestGDriveStorageConnected(TestStorageBase): @@ -63,19 +68,17 @@ class TestGDriveStorageConnected(TestStorageBase): module_name: str = "gdrive_storage" storage: Type[GDriveStorage] - config: dict = {'path_generator': 'url', - 'filename_generator': 'static', - # TODO: replace with real root folder id - 'root_folder_id': "1TVY_oJt95_dmRSEdP9m5zFy7l50TeCSk", - 'oauth_token': None, - 'service_account': 'secrets/service_account.json' - } - + config: dict = { + "path_generator": "url", + "filename_generator": "static", + # TODO: replace with real root folder id + "root_folder_id": "1TVY_oJt95_dmRSEdP9m5zFy7l50TeCSk", + "oauth_token": None, + "service_account": "secrets/service_account.json", + } def test_initialize_with_real_credentials(self): """ Test that the Google Drive service can be initialized with real credentials. """ assert self.storage.service is not None - - diff --git a/tests/storages/test_local_storage.py b/tests/storages/test_local_storage.py index c3581df..1230e3d 100644 --- a/tests/storages/test_local_storage.py +++ b/tests/storages/test_local_storage.py @@ -1,4 +1,3 @@ - import os from pathlib import Path @@ -8,6 +7,7 @@ from auto_archiver.core import Media, Metadata from auto_archiver.modules.local_storage import LocalStorage from auto_archiver.core.consts import SetupError + @pytest.fixture def local_storage(setup_module, tmp_path) -> LocalStorage: save_to = tmp_path / "local_archive" @@ -20,6 +20,7 @@ def local_storage(setup_module, tmp_path) -> LocalStorage: } return setup_module("local_storage", configs) + @pytest.fixture def sample_media(tmp_path) -> Media: """Fixture creating a Media object with temporary source file""" @@ -27,9 +28,11 @@ def sample_media(tmp_path) -> Media: src_file.write_text("test content") return Media(filename=str(src_file)) + def test_too_long_save_path(setup_module): with pytest.raises(SetupError): - setup_module("local_storage", {"save_to": "long"*100}) + setup_module("local_storage", {"save_to": "long" * 100}) + def test_get_cdn_url_relative(local_storage): local_storage.filename_generator = "random" @@ -38,6 +41,7 @@ def test_get_cdn_url_relative(local_storage): expected = os.path.join(local_storage.save_to, media.key) assert local_storage.get_cdn_url(media) == expected + def test_get_cdn_url_absolute(local_storage): local_storage.filename_generator = "random" @@ -47,14 +51,14 @@ def test_get_cdn_url_absolute(local_storage): expected = os.path.abspath(os.path.join(local_storage.save_to, media.key)) assert local_storage.get_cdn_url(media) == expected + def test_upload_file_contents_and_metadata(local_storage, sample_media): local_storage.store(sample_media, "https://example.com", Metadata()) dest = os.path.join(local_storage.save_to, sample_media.key) assert Path(sample_media.filename).read_text() == Path(dest).read_text() + def test_upload_nonexistent_source(local_storage): media = Media(_key="missing.txt", filename="nonexistent.txt") with pytest.raises(FileNotFoundError): local_storage.upload(media) - - diff --git a/tests/storages/test_storage_base.py b/tests/storages/test_storage_base.py index 53dfbd7..730304e 100644 --- a/tests/storages/test_storage_base.py +++ b/tests/storages/test_storage_base.py @@ -6,32 +6,28 @@ from auto_archiver.core.metadata import Metadata, Media from auto_archiver.core.storage import Storage from auto_archiver.core.module import ModuleFactory -class TestStorageBase(object): +class TestStorageBase(object): module_name: str = None config: dict = None @pytest.fixture(autouse=True) def setup_storage(self, setup_module): - assert ( - self.module_name is not None - ), "self.module_name must be set on the subclass" + assert self.module_name is not None, "self.module_name must be set on the subclass" assert self.config is not None, "self.config must be a dict set on the subclass" - self.storage: Type[Storage] = setup_module( - self.module_name, self.config - ) + self.storage: Type[Storage] = setup_module(self.module_name, self.config) class TestBaseStorage(Storage): - name = "test_storage" def get_cdn_url(self, media): return "cdn_url" - + def uploadf(self, file, key, **kwargs): return True + @pytest.fixture def dummy_file(tmp_path): # create dummy.txt file @@ -39,16 +35,18 @@ def dummy_file(tmp_path): dummy_file.write_text("test content") return str(dummy_file) + @pytest.fixture def storage_base(): def _storage_base(config): storage_base = TestBaseStorage() - storage_base.config_setup({TestBaseStorage.name : config}) + storage_base.config_setup({TestBaseStorage.name: config}) storage_base.module_factory = ModuleFactory() return storage_base - + return _storage_base + @pytest.mark.parametrize( "path_generator, filename_generator, url, expected_key", [ @@ -58,11 +56,11 @@ def storage_base(): ("url", "random", "https://example.com/file/", "folder/https-example-com-file/pretend-random.txt"), ("random", "static", "https://example.com/file/", "folder/pretend-random/6ae8a75555209fd6c44157c0.txt"), ("random", "random", "https://example.com/file/", "folder/pretend-random/pretend-random.txt"), - ], ) -def test_storage_name_generation(storage_base, path_generator, filename_generator, url, - expected_key, mocker, tmp_path, dummy_file): +def test_storage_name_generation( + storage_base, path_generator, filename_generator, url, expected_key, mocker, tmp_path, dummy_file +): mock_random = mocker.patch("auto_archiver.core.storage.random_str") mock_random.return_value = "pretend-random" @@ -89,7 +87,24 @@ def test_really_long_name(storage_base, dummy_file): } storage: Storage = storage_base(config) - url = f"https://example.com/{'file'*100}" + url = f"https://example.com/{'file' * 100}" media = Media(filename=dummy_file) storage.set_key(media, url, Metadata()) - assert media.key == f"https-example-com-{'file'*13}/6ae8a75555209fd6c44157c0.txt" \ No newline at end of file + assert media.key == f"https-example-com-{'file' * 13}/6ae8a75555209fd6c44157c0.txt" + + +def test_storage_loads_hash_enricher(storage_base, dummy_file): + """Ensure 'hash_enricher' is properly loaded without an explicit import.""" + config = {"path_generator": "url", "filename_generator": "static"} + storage = storage_base(config) + + url = "https://example.com/file/" + media = Media(filename=dummy_file) + metadata = Metadata() + + try: + storage.set_key(media, url, metadata) + except Exception as e: + pytest.fail(f"Storage failed to dynamically load hash_enricher: {e}") + + assert media.key is not None, "Expected media.key to be set, but it was None" diff --git a/tests/test_config.py b/tests/test_config.py index 75fe515..03b06e7 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -3,39 +3,46 @@ from auto_archiver.core import config from ruamel.yaml.scanner import ScannerError from ruamel.yaml.comments import CommentedMap + def test_return_default_config_for_nonexistent_file(): assert config.read_yaml("nonexistent_file.yaml") == config.EMPTY_CONFIG + def test_return_default_config_for_empty_file(tmp_path): empty_file = tmp_path / "empty_file.yaml" empty_file.write_text("") assert config.read_yaml(empty_file) == config.EMPTY_CONFIG + def test_raise_error_on_invalid_yaml(tmp_path): invalid_yaml = tmp_path / "invalid_yaml.yaml" - invalid_yaml.write_text("key: \"value_without_end_quote") + invalid_yaml.write_text('key: "value_without_end_quote') # make sure it raises ScannerError with pytest.raises(ScannerError): config.read_yaml(invalid_yaml) + def test_write_yaml(tmp_path): yaml_file = tmp_path / "write_yaml.yaml" config.store_yaml(config.EMPTY_CONFIG, yaml_file.as_posix()) assert "steps:\n" in yaml_file.read_text() + def test_round_trip_comments(tmp_path): yaml_file = tmp_path / "round_trip_comments.yaml" with open(yaml_file, "w") as f: - f.write("generic_extractor:\n facebook_cookie: abc # end of line comment\n subtitles: true\n # comments: false\n # livestreams: false\n list_type:\n - value1\n - value2") + f.write( + "generic_extractor:\n facebook_cookie: abc # end of line comment\n subtitles: true\n # comments: false\n # livestreams: false\n list_type:\n - value1\n - value2" + ) loaded = config.read_yaml(yaml_file) # check the comments are preserved - assert loaded['generic_extractor']['facebook_cookie'] == "abc" - assert loaded['generic_extractor'].ca.items['facebook_cookie'][2].value == "# end of line comment\n" + assert loaded["generic_extractor"]["facebook_cookie"] == "abc" + assert loaded["generic_extractor"].ca.items["facebook_cookie"][2].value == "# end of line comment\n" # add some more items to my_settings - loaded['generic_extractor']['list_type'].append("bellingcat") + loaded["generic_extractor"]["list_type"].append("bellingcat") config.store_yaml(loaded, yaml_file.as_posix()) assert "# comments: false" in yaml_file.read_text() @@ -43,14 +50,17 @@ def test_round_trip_comments(tmp_path): assert "abc # end of line comment" in yaml_file.read_text() assert "- value2\n - bellingcat" in yaml_file.read_text() + def test_merge_dicts(): yaml_dict = config.EMPTY_CONFIG - yaml_dict['settings'] = CommentedMap(**{ + yaml_dict["settings"] = CommentedMap( + **{ "key1": ["a"], "key2": "old_value", "key3": ["a", "b", "c"], "key5": "value5", - }) + } + ) dotdict = { "settings.key1": ["b", "c"], @@ -67,15 +77,16 @@ def test_merge_dicts(): def test_check_types(): - assert config.is_list_type([]) == True - assert config.is_list_type(()) == True - assert config.is_list_type(set()) == True - assert config.is_list_type({}) == False - assert config.is_list_type("") == False - assert config.is_dict_type({}) == True - assert config.is_dict_type(CommentedMap()) == True - assert config.is_dict_type([]) == False - assert config.is_dict_type("") == False + assert config.is_list_type([]) is True + assert config.is_list_type(()) is True + assert config.is_list_type(set()) is True + assert config.is_list_type({}) is False + assert config.is_list_type("") is False + assert config.is_dict_type({}) is True + assert config.is_dict_type(CommentedMap()) is True + assert config.is_dict_type([]) is False + assert config.is_dict_type("") is False + def test_from_dot_notation(): dotdict = { @@ -88,16 +99,17 @@ def test_from_dot_notation(): assert normal_dict["settings"]["key2"] == "new_value" assert normal_dict["settings"]["key3"]["key4"] == "value" + def test_to_dot_notation(): yaml_dict = config.EMPTY_CONFIG - yaml_dict['settings'] = { + yaml_dict["settings"] = { "key1": ["a", "b", "c"], "key2": "new_value", "key3": { "key4": "value", - } + }, } dotdict = config.to_dot_notation(yaml_dict) assert dotdict["settings.key1"] == ["a", "b", "c"] assert dotdict["settings.key2"] == "new_value" - assert dotdict["settings.key3.key4"] == "value" \ No newline at end of file + assert dotdict["settings.key3.key4"] == "value" diff --git a/tests/test_implementation.py b/tests/test_implementation.py index 51e9d79..e52a8d8 100644 --- a/tests/test_implementation.py +++ b/tests/test_implementation.py @@ -10,21 +10,23 @@ def orchestration_file_path(tmp_path): folder.mkdir(exist_ok=True) return (folder / "example_orch.yaml").as_posix() + @pytest.fixture def orchestration_file(orchestration_file_path): - def _orchestration_file(content=''): + def _orchestration_file(content=""): with open(orchestration_file_path, "w") as f: f.write(content) return orchestration_file_path - + return _orchestration_file + @pytest.fixture def autoarchiver(tmp_path, monkeypatch, request): def _autoarchiver(args=[]): - def cleanup(): from loguru import logger + if not logger._core.handlers.get(0): logger._core.handlers_count = 0 logger.add(sys.stderr) @@ -44,9 +46,9 @@ def autoarchiver(tmp_path, monkeypatch, request): def test_run_auto_archiver_no_args(caplog, autoarchiver): with pytest.raises(SystemExit): autoarchiver() - assert "provide at least one URL via the command line, or set up an alternative feeder" in caplog.text + def test_run_auto_archiver_invalid_file(caplog, autoarchiver): # exec 'auto-archiver' on the command lin with pytest.raises(SystemExit): @@ -54,6 +56,7 @@ def test_run_auto_archiver_invalid_file(caplog, autoarchiver): assert "Make sure the file exists and try again, or run without th" in caplog.text + def test_run_auto_archiver_empty_file(caplog, autoarchiver, orchestration_file): # create a valid (empty) orchestration file path = orchestration_file(content="") @@ -64,6 +67,7 @@ def test_run_auto_archiver_empty_file(caplog, autoarchiver, orchestration_file): # should treat an empty file as if there is no file at all assert " No URLs provided. Please provide at least one URL via the com" in caplog.text + def test_call_autoarchiver_main(caplog, monkeypatch, tmp_path): from auto_archiver.__main__ import main @@ -75,4 +79,4 @@ def test_call_autoarchiver_main(caplog, monkeypatch, tmp_path): with pytest.raises(SystemExit): main() - assert "No URLs provided. Please provide at least one" in caplog.text \ No newline at end of file + assert "No URLs provided. Please provide at least one" in caplog.text diff --git a/tests/test_metadata.py b/tests/test_metadata.py index e1f7797..e838979 100644 --- a/tests/test_metadata.py +++ b/tests/test_metadata.py @@ -62,18 +62,8 @@ def test_simple_merge(basic_metadata): def test_left_merge(): - left = ( - Metadata() - .set("tags", ["a"]) - .set("stats", {"views": 10}) - .set("status", "success") - ) - right = ( - Metadata() - .set("tags", ["b"]) - .set("stats", {"likes": 5}) - .set("status", "no archiver") - ) + left = Metadata().set("tags", ["a"]).set("stats", {"views": 10}).set("status", "success") + right = Metadata().set("tags", ["b"]).set("stats", {"likes": 5}).set("status", "no archiver") left.merge(right, overwrite_left=True) assert left.get("status") == "no archiver" @@ -120,6 +110,7 @@ def test_is_empty(): def test_store(): pass + # Test Media operations @@ -176,6 +167,7 @@ def test_choose_most_complete(): res = Metadata.choose_most_complete([m_more, m_less]) assert res.metadata.get("title") == "Title 1" + def test_choose_most_complete_from_pickles(unpickle): # test most complete from pickles before and after an enricher has run # Only compares length of media, not the actual media diff --git a/tests/test_modules.py b/tests/test_modules.py index 7a2b14d..f672ca6 100644 --- a/tests/test_modules.py +++ b/tests/test_modules.py @@ -1,25 +1,25 @@ -import sys import pytest from auto_archiver.core.module import ModuleFactory, LazyBaseModule from auto_archiver.core.base_module import BaseModule + @pytest.fixture def example_module(): import auto_archiver module_factory = ModuleFactory() - - previous_path = auto_archiver.modules.__path__ + # previous_path = auto_archiver.modules.__path__ auto_archiver.modules.__path__.append("tests/data/test_modules/") - return module_factory.get_module_lazy("example_module") + def test_get_module_lazy(example_module): assert example_module.name == "example_module" assert example_module.display_name == "Example Module" assert example_module.manifest is not None + def test_python_dependency_check(example_module): # example_module requires loguru, which is not installed # monkey patch the manifest to include a nonexistnet dependency @@ -30,11 +30,13 @@ def test_python_dependency_check(example_module): assert load_error.value.code == 1 + def test_binary_dependency_check(example_module): # example_module requires ffmpeg, which is not installed # monkey patch the manifest to include a nonexistnet dependency example_module.manifest["dependencies"]["binary"] = ["does_not_exist"] + def test_module_dependency_check_loads_module(example_module): # example_module requires cli_feeder, which is not installed # monkey patch the manifest to include a nonexistnet dependency @@ -49,19 +51,20 @@ def test_module_dependency_check_loads_module(example_module): assert module_factory._lazy_modules["hash_enricher"] is not None assert module_factory._lazy_modules["hash_enricher"]._instance is not None -def test_load_module(example_module): +def test_load_module(example_module): # setup the module, and check that config is set to the default values loaded_module = example_module.load({}) assert loaded_module is not None assert isinstance(loaded_module, BaseModule) assert loaded_module.name == "example_module" assert loaded_module.display_name == "Example Module" - assert loaded_module.config["example_module"] == {"csv_file" : "db.csv"} + assert loaded_module.config["example_module"] == {"csv_file": "db.csv"} # check that the vlaue is set on the module itself assert loaded_module.csv_file == "db.csv" + @pytest.mark.parametrize("module_name", ["local_storage", "generic_extractor", "html_formatter", "csv_db"]) def test_load_modules(module_name): # test that specific modules can be loaded @@ -78,6 +81,8 @@ def test_load_modules(module_name): # check that default settings are applied default_config = module.configs assert loaded_module.name in loaded_module.config.keys() + defaults = {k: v.get("default") for k, v in default_config.items()} + assert loaded_module.config[module_name] == defaults @pytest.mark.parametrize("module_name", ["local_storage", "generic_extractor", "html_formatter", "csv_db"]) @@ -96,5 +101,3 @@ def test_lazy_base_module(module_name): assert len(lazy_module.configs) > 0 assert len(lazy_module.description) > 0 assert len(lazy_module.version) > 0 - - diff --git a/tests/test_orchestrator.py b/tests/test_orchestrator.py index 72f4949..326b93d 100644 --- a/tests/test_orchestrator.py +++ b/tests/test_orchestrator.py @@ -1,5 +1,4 @@ import pytest -import sys from argparse import ArgumentParser, ArgumentTypeError from auto_archiver.core.orchestrator import ArchivingOrchestrator from auto_archiver.version import __version__ @@ -9,51 +8,65 @@ from auto_archiver.core import Metadata TEST_ORCHESTRATION = "tests/data/test_orchestration.yaml" TEST_MODULES = "tests/data/test_modules/" + @pytest.fixture def test_args(): - return ["--config", TEST_ORCHESTRATION, - "--module_paths", TEST_MODULES, - "--example_module.required_field", "some_value"] # just set this for normal testing, we will remove it later + return [ + "--config", + TEST_ORCHESTRATION, + "--module_paths", + TEST_MODULES, + "--example_module.required_field", + "some_value", + ] # just set this for normal testing, we will remove it later + @pytest.fixture def orchestrator(): return ArchivingOrchestrator() + @pytest.fixture def basic_parser(orchestrator) -> ArgumentParser: return orchestrator.setup_basic_parser() + def test_setup_orchestrator(orchestrator): assert orchestrator is not None + def test_parse_config(): pass + def test_parse_basic(basic_parser): args = basic_parser.parse_args(["--config", TEST_ORCHESTRATION]) assert args.config_file == TEST_ORCHESTRATION + @pytest.mark.parametrize("mode", ["simple", "full"]) def test_mode(basic_parser, mode): args = basic_parser.parse_args(["--mode", mode]) assert args.mode == mode + def test_mode_invalid(basic_parser, capsys): with pytest.raises(SystemExit) as exit_error: basic_parser.parse_args(["--mode", "invalid"]) assert exit_error.value.code == 2 assert "invalid choice" in capsys.readouterr().err + def test_version(basic_parser, capsys): with pytest.raises(SystemExit) as exit_error: basic_parser.parse_args(["--version"]) assert exit_error.value.code == 0 assert capsys.readouterr().out == f"{__version__}\n" -def test_help(orchestrator, basic_parser, capsys): +def test_help(orchestrator, basic_parser, capsys): args = basic_parser.parse_args(["--help"]) - assert args.help == True + assert args.help is True # test the show_help() on orchestrator with pytest.raises(SystemExit) as exit_error: @@ -83,14 +96,17 @@ def test_help(orchestrator, basic_parser, capsys): def test_add_custom_modules_path(orchestrator, test_args): orchestrator.setup_config(test_args) - + import auto_archiver + assert "tests/data/test_modules/" in auto_archiver.modules.__path__ -def test_add_custom_modules_path_invalid(orchestrator, caplog, test_args): - orchestrator.setup_config(test_args + # we still need to load the real path to get the example_module - ["--module_paths", "tests/data/invalid_test_modules/"]) +def test_add_custom_modules_path_invalid(orchestrator, caplog, test_args): + orchestrator.setup_config( + test_args # we still need to load the real path to get the example_module + + ["--module_paths", "tests/data/invalid_test_modules/"] + ) assert caplog.records[0].message == "Path 'tests/data/invalid_test_modules/' does not exist. Skipping..." @@ -99,16 +115,16 @@ def test_check_required_values(orchestrator, caplog, test_args): # drop the example_module.required_field from the test_args test_args = test_args[:-2] - with pytest.raises(SystemExit) as exit_error: - config = orchestrator.setup_config(test_args) + with pytest.raises(SystemExit): + orchestrator.setup_config(test_args) assert caplog.records[1].message == "the following arguments are required: --example_module.required_field" -def test_get_required_values_from_config(orchestrator, test_args, tmp_path): +def test_get_required_values_from_config(orchestrator, test_args, tmp_path): # load the default example yaml, add a required field, then run the orchestrator test_yaml = read_yaml(TEST_ORCHESTRATION) - test_yaml['example_module'] = {'required_field': 'some_value'} + test_yaml["example_module"] = {"required_field": "some_value"} # write it to a temp file tmp_file = (tmp_path / "temp_config.yaml").as_posix() store_yaml(test_yaml, tmp_file) @@ -117,27 +133,42 @@ def test_get_required_values_from_config(orchestrator, test_args, tmp_path): config = orchestrator.setup_config(["--config", tmp_file, "--module_paths", TEST_MODULES]) assert config is not None -def test_load_authentication_string(orchestrator, test_args): - config = orchestrator.setup_config(test_args + ["--authentication", '{"facebook.com": {"username": "my_username", "password": "my_password"}}']) - assert config['authentication'] == {"facebook.com": {"username": "my_username", "password": "my_password"}} +def test_load_authentication_string(orchestrator, test_args): + config = orchestrator.setup_config( + test_args + ["--authentication", '{"facebook.com": {"username": "my_username", "password": "my_password"}}'] + ) + assert config["authentication"] == {"facebook.com": {"username": "my_username", "password": "my_password"}} + def test_load_authentication_string_concat_site(orchestrator, test_args): - config = orchestrator.setup_config(test_args + ["--authentication", '{"x.com,twitter.com": {"api_key": "my_key"}}']) - assert config['authentication'] == {"x.com": {"api_key": "my_key"}, - "twitter.com": {"api_key": "my_key"}} + assert config["authentication"] == {"x.com": {"api_key": "my_key"}, "twitter.com": {"api_key": "my_key"}} + def test_load_invalid_authentication_string(orchestrator, test_args): with pytest.raises(ArgumentTypeError): - orchestrator.setup_config(test_args + ["--authentication", "{\''invalid_json"]) + orchestrator.setup_config(test_args + ["--authentication", "{''invalid_json"]) + def test_load_authentication_invalid_dict(orchestrator, test_args): with pytest.raises(ArgumentTypeError): orchestrator.setup_config(test_args + ["--authentication", "[true, false]"]) + def test_load_modules_from_commandline(orchestrator, test_args): - args = test_args + ["--feeders", "example_module", "--extractors", "example_module", "--databases", "example_module", "--enrichers", "example_module", "--formatters", "example_module"] + args = test_args + [ + "--feeders", + "example_module", + "--extractors", + "example_module", + "--databases", + "example_module", + "--enrichers", + "example_module", + "--formatters", + "example_module", + ] orchestrator.setup(args) @@ -153,27 +184,37 @@ def test_load_modules_from_commandline(orchestrator, test_args): assert orchestrator.enrichers[0].name == "example_module" assert orchestrator.formatters[0].name == "example_module" + def test_load_settings_for_module_from_commandline(orchestrator, test_args): - args = test_args + ["--feeders", "gsheet_feeder_db", "--gsheet_feeder_db.sheet_id", "123", "--gsheet_feeder_db.service_account", "tests/data/test_service_account.json"] + args = test_args + [ + "--feeders", + "gsheet_feeder_db", + "--gsheet_feeder_db.sheet_id", + "123", + "--gsheet_feeder_db.service_account", + "tests/data/test_service_account.json", + ] orchestrator.setup(args) assert len(orchestrator.feeders) == 1 assert orchestrator.feeders[0].name == "gsheet_feeder_db" - assert orchestrator.config['gsheet_feeder_db']['sheet_id'] == "123" + assert orchestrator.config["gsheet_feeder_db"]["sheet_id"] == "123" def test_multiple_orchestrator(test_args): - - o1_args = test_args + ["--feeders", "gsheet_feeder_db", "--gsheet_feeder_db.service_account", "tests/data/test_service_account.json"] + o1_args = test_args + [ + "--feeders", + "gsheet_feeder_db", + "--gsheet_feeder_db.service_account", + "tests/data/test_service_account.json", + ] o1 = ArchivingOrchestrator() - with pytest.raises(ValueError) as exit_error: + with pytest.raises(ValueError): # this should fail because the gsheet_feeder_db requires a sheet_id / sheet o1.setup(o1_args) - - o2_args = test_args + ["--feeders", "example_module"] o2 = ArchivingOrchestrator() o2.setup(o2_args) @@ -182,4 +223,4 @@ def test_multiple_orchestrator(test_args): output: Metadata = list(o2.feed()) assert len(output) == 1 - assert output[0].get_url() == "https://example.com" \ No newline at end of file + assert output[0].get_url() == "https://example.com" diff --git a/tests/utils/test_misc.py b/tests/utils/test_misc.py index 0023077..844d3d2 100644 --- a/tests/utils/test_misc.py +++ b/tests/utils/test_misc.py @@ -14,7 +14,7 @@ from auto_archiver.utils.misc import ( update_nested_dict, calculate_file_hash, random_str, - get_timestamp + get_timestamp, ) @@ -38,40 +38,46 @@ class TestDirectoryUtils: mkdir_if_not_exists(existing_dir) assert existing_dir.exists() + class TestURLExpansion: - @pytest.mark.parametrize("input_url,expected", [ - ("https://example.com", "https://example.com"), - ("https://t.co/test", "https://expanded.url") - ]) + @pytest.mark.parametrize( + "input_url,expected", + [("https://example.com", "https://example.com"), ("https://t.co/test", "https://expanded.url")], + ) def test_expand_url(self, input_url, expected, mocker): mock_response = mocker.Mock() mock_response.url = "https://expanded.url" - mocker.patch('requests.get', return_value=mock_response) + mocker.patch("requests.get", return_value=mock_response) result = expand_url(input_url) assert result == expected def test_expand_url_handles_errors(self, caplog, mocker): - mocker.patch('requests.get', side_effect=Exception("Connection error")) + mocker.patch("requests.get", side_effect=Exception("Connection error")) url = "https://t.co/error" result = expand_url(url) assert result == url assert f"Failed to expand url {url}" in caplog.text + class TestAttributeHandling: class Sample: exists = "value" none = None - @pytest.mark.parametrize("obj,attr,default,expected", [ - (Sample(), "exists", "default", "value"), - (Sample(), "none", "default", "default"), - (Sample(), "missing", "default", "default"), - (None, "anything", "fallback", "fallback"), - ]) + @pytest.mark.parametrize( + "obj,attr,default,expected", + [ + (Sample(), "exists", "default", "value"), + (Sample(), "none", "default", "default"), + (Sample(), "missing", "default", "default"), + (None, "anything", "fallback", "fallback"), + ], + ) def test_getattr_or(self, obj, attr, default, expected): # Test gets attribute or returns a default value assert getattr_or(obj, attr, default) == expected + class TestDateTimeHandling: def test_datetime_encoder(self, sample_datetime): result = json.dumps({"dt": sample_datetime}, cls=DateTimeEncoder) @@ -83,11 +89,14 @@ class TestDateTimeHandling: result = dump_payload(payload) assert str(sample_datetime) in result - @pytest.mark.parametrize("dt_str,fmt,expected", [ - ("2023-01-01 12:00:00+00:00", None, datetime(2023, 1, 1, 12, 0, tzinfo=timezone.utc)), - ("20230101 120000", "%Y%m%d %H%M%S", datetime(2023, 1, 1, 12, 0)), - ("invalid", None, None), - ]) + @pytest.mark.parametrize( + "dt_str,fmt,expected", + [ + ("2023-01-01 12:00:00+00:00", None, datetime(2023, 1, 1, 12, 0, tzinfo=timezone.utc)), + ("20230101 120000", "%Y%m%d %H%M%S", datetime(2023, 1, 1, 12, 0)), + ("invalid", None, None), + ], + ) def test_datetime_from_string(self, dt_str, fmt, expected): result = get_datetime_from_str(dt_str, fmt) if expected is None: @@ -95,16 +104,21 @@ class TestDateTimeHandling: else: assert result == expected.replace(tzinfo=result.tzinfo) + class TestDictUtils: - @pytest.mark.parametrize("original,update,expected", [ - ({"a": 1}, {"b": 2}, {"a": 1, "b": 2}), - ({"nested": {"a": 1}}, {"nested": {"b": 2}}, {"nested": {"a": 1, "b": 2}}), - ({"a": {"b": {"c": 1}}}, {"a": {"b": {"c": 2}}}, {"a": {"b": {"c": 2}}}), - ]) + @pytest.mark.parametrize( + "original,update,expected", + [ + ({"a": 1}, {"b": 2}, {"a": 1, "b": 2}), + ({"nested": {"a": 1}}, {"nested": {"b": 2}}, {"nested": {"a": 1, "b": 2}}), + ({"a": {"b": {"c": 1}}}, {"a": {"b": {"c": 2}}}, {"a": {"b": {"c": 2}}}), + ], + ) def test_update_nested_dict(self, original, update, expected): update_nested_dict(original, update) assert original == expected + class TestHashingUtils: def test_file_hashing(self, sample_file): expected = hashlib.sha256(b"test content").hexdigest() @@ -118,6 +132,7 @@ class TestHashingUtils: expected = hashlib.sha256(content).hexdigest() assert calculate_file_hash(str(file_path)) == expected + class TestMiscUtils: def test_random_str_length(self): for length in [8, 16, 32]: @@ -131,14 +146,17 @@ class TestMiscUtils: def test_random_str_uniqueness(self): assert random_str() != random_str() - @pytest.mark.parametrize("ts_input,utc,iso,expected_type", [ - (datetime.now(), True, True, str), - ("2023-01-01T12:00:00+00:00", False, False, datetime), - (1672574400, True, True, str), - ]) + @pytest.mark.parametrize( + "ts_input,utc,iso,expected_type", + [ + (datetime.now(), True, True, str), + ("2023-01-01T12:00:00+00:00", False, False, datetime), + (1672574400, True, True, str), + ], + ) def test_timestamp_parsing(self, ts_input, utc, iso, expected_type): result = get_timestamp(ts_input, utc=utc, iso=iso) assert isinstance(result, expected_type) def test_invalid_timestamp_returns_none(self): - assert get_timestamp("invalid-date") is None \ No newline at end of file + assert get_timestamp("invalid-date") is None