diff --git a/src/auto_archiver/core/validators.py b/src/auto_archiver/core/validators.py index 765a02b..decd9f1 100644 --- a/src/auto_archiver/core/validators.py +++ b/src/auto_archiver/core/validators.py @@ -4,12 +4,6 @@ import argparse import json -def example_validator(value): - if "example" not in value: - raise argparse.ArgumentTypeError(f"{value} is not a valid value for this argument") - return value - - def positive_number(value): if value < 0: raise argparse.ArgumentTypeError(f"{value} is not a positive number") diff --git a/src/auto_archiver/modules/telethon_extractor/__manifest__.py b/src/auto_archiver/modules/telethon_extractor/__manifest__.py index 150b62c..e29a9b0 100644 --- a/src/auto_archiver/modules/telethon_extractor/__manifest__.py +++ b/src/auto_archiver/modules/telethon_extractor/__manifest__.py @@ -19,7 +19,7 @@ }, "session_file": { "default": "secrets/anon", - "help": "optional, records the telegram login session for future usage, '.session' will be appended to the provided value.", + "help": "Path of the file to save the telegram login session for future usage, '.session' will be appended to the provided path.", }, "join_channels": { "default": True, diff --git a/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py b/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py index 434d97c..9acec30 100644 --- a/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py +++ b/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py @@ -1,4 +1,9 @@ +import os import shutil +import re +import time +from datetime import date + from telethon.sync import TelegramClient from telethon.errors import ChannelInvalidError from telethon.tl.functions.messages import ImportChatInviteRequest @@ -8,11 +13,9 @@ from telethon.errors.rpcerrorlist import ( InviteRequestSentError, InviteHashExpiredError, ) -from loguru import logger + from tqdm import tqdm -import re -import time -import os +from loguru import logger from auto_archiver.core import Extractor from auto_archiver.core import Metadata, Media @@ -32,13 +35,22 @@ class TelethonExtractor(Extractor): logger.info(f"SETUP {self.name} checking login...") # in case the user already added '.session' to the session_file - self.session_file = self.session_file.removesuffix(".session") + base_session_name = self.session_file.removesuffix(".session") + base_session_filepath = f"{base_session_name}.session" + + if self.session_file and not os.path.exists(base_session_filepath): + logger.warning( + f"SETUP - Session file {base_session_filepath} does not exist for {self.name}, creating an empty one." + ) + with open(base_session_filepath, "w") as f: + f.write("") # make a copy of the session that is used exclusively with this archiver instance - old_session_file = f"{self.session_file}.session" - self.session_file = os.path.join("secrets/", f"telethon-{time.strftime('%Y-%m-%d')}{random_str(8)}") - logger.debug(f"Making a copy of the session file {old_session_file} to {self.session_file}.session") - shutil.copy(old_session_file, f"{self.session_file}.session") + self.session_file = os.path.join( + os.path.dirname(base_session_filepath), f"telethon-{date.today().strftime('%Y-%m-%d')}{random_str(8)}" + ) + logger.debug(f"Making a copy of the session file {base_session_filepath} to {self.session_file}.session") + shutil.copy(base_session_filepath, f"{self.session_file}.session") # initiate the client self.client = TelegramClient(self.session_file, self.api_id, self.api_hash) diff --git a/tests/conftest.py b/tests/conftest.py index 44d8058..b686319 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -20,6 +20,14 @@ from auto_archiver.core.module import ModuleFactory TESTS_TO_RUN_LAST = ["test_twitter_api_archiver"] +@pytest.fixture +def get_lazy_module(): + def _get_lazy_module(module_name): + return ModuleFactory().get_module_lazy(module_name) + + return _get_lazy_module + + @pytest.fixture def setup_module(request): def _setup_module(module_name, config=None): diff --git a/tests/extractors/test_telethon_extractor.py b/tests/extractors/test_telethon_extractor.py new file mode 100644 index 0000000..23ba338 --- /dev/null +++ b/tests/extractors/test_telethon_extractor.py @@ -0,0 +1,38 @@ +import os +from datetime import date + +import pytest + +from .test_extractor_base import TestExtractorBase +from auto_archiver.modules.telethon_extractor import TelethonExtractor + + +class TestTelethonExtractor(TestExtractorBase): + extractor_module = "telethon_extractor" + extractor: TelethonExtractor + config = { + "api_id": 123, + "api_hash": "ABC", + } + + @pytest.fixture(autouse=True) + def mock_client_setup(self, mocker): + mocker.patch("telethon.client.auth.AuthMethods.start") + + def test_setup_fails_clear_session_file(self, get_lazy_module, tmp_path, mocker): + start = mocker.patch("telethon.client.auth.AuthMethods.start") + start.side_effect = Exception("Test exception") + + # make sure the default setup file is created + session_file = tmp_path / "test.session" + + lazy_module = get_lazy_module("telethon_extractor") + + with pytest.raises(Exception): + lazy_module.load( + {"telethon_extractor": {"session_file": str(session_file), "api_id": 123, "api_hash": "ABC"}} + ) + + assert session_file.exists() + assert f"telethon-{date.today().strftime('%Y-%m-%d')}" in lazy_module._instance.session_file + assert os.path.exists(lazy_module._instance.session_file + ".session")