Telethon unit tests + tidyup

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")

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,

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)

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):

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")