extracting qbittorrent logic into separate class

media_manager/torrent/download_clients/abstractDownloadClient.py

@@ -0,0 +1,41 @@
+from abc import ABC, abstractmethod
+
+from media_manager.indexer.schemas import IndexerQueryResult
+from media_manager.torrent.schemas import TorrentId, TorrentStatus, Torrent
+
+
+class AbstractDownloadClient(ABC):
+    """
+    Abstract base class for download clients.
+    Defines the interface that all download clients must implement.
+    """
+
+    @abstractmethod
+    def download_torrent(self, torrent: IndexerQueryResult) -> Torrent:
+        """
+        Add a torrent to the download client and return the torrent object.
+
+        :param torrent: The indexer query result of the torrent file to download.
+        :return: The torrent object with calculated hash and initial status.
+        """
+        pass
+
+    @abstractmethod
+    def remove_torrent(self, torrent: Torrent, delete_data: bool = False) -> None:
+        """
+        Remove a torrent from the download client.
+
+        :param torrent: The torrent to remove.
+        :param delete_data: Whether to delete the downloaded data.
+        """
+        pass
+
+    @abstractmethod
+    def get_torrent_status(self, torrent: Torrent) -> TorrentStatus:
+        """
+        Get the status of a specific torrent.
+
+        :param torrent: The torrent to get the status of.
+        :return: The status of the torrent.
+        """
+        pass

media_manager/torrent/download_clients/qbittorrent.py

@@ -0,0 +1,197 @@
+import hashlib
+import logging
+
+import bencoder
+import qbittorrentapi
+import requests
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+from media_manager.config import BasicConfig
+from media_manager.indexer.schemas import IndexerQueryResult
+from media_manager.torrent.download_clients.abstractDownloadClient import AbstractDownloadClient
+from media_manager.torrent.schemas import TorrentId, TorrentStatus, Torrent
+
+log = logging.getLogger(__name__)
+
+
+class QbittorrentConfig(BaseSettings):
+    model_config = SettingsConfigDict(env_prefix="QBITTORRENT_")
+    host: str = "localhost"
+    port: int = 8080
+    username: str = "admin"
+    password: str = "admin"
+
+
+class QbittorrentDownloadClient(AbstractDownloadClient):
+    DOWNLOADING_STATE = (
+        "allocating",
+        "downloading",
+        "metaDL",
+        "pausedDL",
+        "queuedDL",
+        "stalledDL",
+        "checkingDL",
+        "forcedDL",
+        "moving",
+    )
+    FINISHED_STATE = (
+        "uploading",
+        "pausedUP",
+        "queuedUP",
+        "stalledUP",
+        "checkingUP",
+        "forcedUP",
+    )
+    ERROR_STATE = ("missingFiles", "error", "checkingResumeData")
+    UNKNOWN_STATE = ("unknown",)
+
+    def __init__(self):
+        self.config = QbittorrentConfig()
+        self.api_client = qbittorrentapi.Client(**self.config.model_dump())
+        try:
+            self.api_client.auth_log_in()
+            log.info("Successfully logged into qbittorrent")
+        except Exception as e:
+            log.error(f"Failed to log into qbittorrent: {e}")
+            raise
+        finally:
+            self.api_client.auth_log_out()
+
+    def download_torrent(self, indexer_result: IndexerQueryResult) -> Torrent:
+        """
+        Add a torrent to the download client and return the torrent object.
+
+        :param indexer_result: The indexer query result of the torrent file to download.
+        :return: The torrent object with calculated hash and initial status.
+        """
+        log.info(f"Attempting to download torrent: {indexer_result.title}")
+
+        torrent_filepath = BasicConfig().torrent_directory / f"{indexer_result.title}.torrent"
+
+        if torrent_filepath.exists():
+            log.warning(f"Torrent already exists: {torrent_filepath}")
+            # Calculate hash from existing file
+            with open(torrent_filepath, "rb") as file:
+                content = file.read()
+                decoded_content = bencoder.decode(content)
+                torrent_hash = hashlib.sha1(bencoder.encode(decoded_content[b"info"])).hexdigest()
+        else:
+            # Download the torrent file
+            with open(torrent_filepath, "wb") as file:
+                content = requests.get(str(indexer_result.download_url)).content
+                file.write(content)
+
+            # Calculate hash and add to qBittorrent
+            with open(torrent_filepath, "rb") as file:
+                content = file.read()
+                try:
+                    decoded_content = bencoder.decode(content)
+                except Exception as e:
+                    log.error(f"Failed to decode torrent file: {e}")
+                    raise e
+
+                torrent_hash = hashlib.sha1(
+                    bencoder.encode(decoded_content[b"info"])
+                ).hexdigest()
+
+                try:
+                    self.api_client.auth_log_in()
+                    answer = self.api_client.torrents_add(
+                        category="MediaManager", torrent_files=content, save_path=indexer_result.title
+                    )
+                finally:
+                    self.api_client.auth_log_out()
+
+            if answer != "Ok.":
+                log.error(f"Failed to download torrent. API response: {answer}")
+                raise RuntimeError(
+                    f"Failed to download torrent, API-Answer isn't 'Ok.'; API Answer: {answer}"
+                )
+
+        log.info(f"Successfully processed torrent: {indexer_result.title}")
+
+        # Create and return torrent object
+        torrent = Torrent(
+            status=TorrentStatus.unknown,
+            title=indexer_result.title,
+            quality=indexer_result.quality,
+            imported=False,
+            hash=torrent_hash,
+        )
+
+        # Get initial status from download client
+        torrent.status = self.get_torrent_status(torrent)
+
+        return torrent
+
+    def remove_torrent(self, torrent: Torrent, delete_data: bool = False) -> None:
+        """
+        Remove a torrent from the download client.
+
+        :param torrent: The torrent to remove.
+        :param delete_data: Whether to delete the downloaded data.
+        """
+        log.info(f"Removing torrent: {torrent.title}")
+        try:
+            self.api_client.auth_log_in()
+            self.api_client.torrents_delete(torrent_hashes=torrent.hash, delete_files=delete_data)
+        finally:
+            self.api_client.auth_log_out()
+
+    def get_torrent_status(self, torrent: Torrent) -> TorrentStatus:
+        """
+        Get the status of a specific torrent.
+
+        :param torrent: The torrent to get the status of.
+        :return: The status of the torrent.
+        """
+        log.info(f"Fetching status for torrent: {torrent.title}")
+        try:
+            self.api_client.auth_log_in()
+            info = self.api_client.torrents_info(torrent_hashes=torrent.hash)
+        finally:
+            self.api_client.auth_log_out()
+
+        if not info:
+            log.warning(f"No information found for torrent: {torrent.id}")
+            return TorrentStatus.unknown
+        else:
+            state: str = info[0]["state"]
+            log.info(f"Torrent {torrent.id} is in state: {state}")
+
+            if state in self.DOWNLOADING_STATE:
+                return TorrentStatus.downloading
+            elif state in self.FINISHED_STATE:
+                return TorrentStatus.finished
+            elif state in self.ERROR_STATE:
+                return TorrentStatus.error
+            elif state in self.UNKNOWN_STATE:
+                return TorrentStatus.unknown
+            else:
+                return TorrentStatus.error
+
+    def pause_torrent(self, torrent: Torrent) -> None:
+        """
+        Pause a torrent download.
+
+        :param torrent: The torrent to pause.
+        """
+        log.info(f"Pausing torrent: {torrent.title}")
+        try:
+            self.api_client.auth_log_in()
+            self.api_client.torrents_pause(torrent_hashes=torrent.hash)
+        finally:
+            self.api_client.auth_log_out()
+
+    def resume_torrent(self, torrent: Torrent) -> None:
+        """
+        Resume a torrent download.
+
+        :param torrent: The torrent to resume.
+        """
+        log.info(f"Resuming torrent: {torrent.title}")
+        try:
+            self.api_client.auth_log_in()
+            self.api_client.torrents_resume(torrent_hashes=torrent.hash)
+        finally:
+            self.api_client.auth_log_out()

media_manager/torrent/service.py

@@ -1,13 +1,8 @@
-import hashlib
 import logging
 
-import bencoder
-import qbittorrentapi
-import requests
-from pydantic_settings import BaseSettings, SettingsConfigDict
-
 from media_manager.config import BasicConfig
 from media_manager.indexer.schemas import IndexerQueryResult
+from media_manager.torrent.download_clients.qbittorrent import QbittorrentDownloadClient
 from media_manager.torrent.repository import TorrentRepository
 from media_manager.torrent.schemas import Torrent, TorrentStatus, TorrentId
 from media_manager.tv.schemas import SeasonFile, Show
@@ -16,48 +11,10 @@ from media_manager.movies.schemas import Movie
 log = logging.getLogger(__name__)
 
 
-class TorrentServiceConfig(BaseSettings):
-    model_config = SettingsConfigDict(env_prefix="QBITTORRENT_")
-    host: str = "localhost"
-    port: int = 8080
-    username: str = "admin"
-    password: str = "admin"
-
-
 class TorrentService:
-    DOWNLOADING_STATE = (
-        "allocating",
-        "downloading",
-        "metaDL",
-        "pausedDL",
-        "queuedDL",
-        "stalledDL",
-        "checkingDL",
-        "forcedDL",
-        "moving",
-    )
-    FINISHED_STATE = (
-        "uploading",
-        "pausedUP",
-        "queuedUP",
-        "stalledUP",
-        "checkingUP",
-        "forcedUP",
-    )
-    ERROR_STATE = ("missingFiles", "error", "checkingResumeData")
-    UNKNOWN_STATE = ("unknown",)
-    api_client = qbittorrentapi.Client(**TorrentServiceConfig().model_dump())
-
-    def __init__(self, torrent_repository: TorrentRepository):
-        try:
-            self.api_client.auth_log_in()
-            log.info("Successfully logged into qbittorrent")
-            self.torrent_repository = torrent_repository
-        except Exception as e:
-            log.error(f"Failed to log into qbittorrent: {e}")
-            raise
-        finally:
-            self.api_client.auth_log_out()
+    def __init__(self, torrent_repository: TorrentRepository, download_client: QbittorrentDownloadClient = None):
+        self.torrent_repository = torrent_repository
+        self.download_client = download_client or QbittorrentDownloadClient()
 
     def get_season_files_of_torrent(self, torrent: Torrent) -> list[SeasonFile]:
         """
@@ -87,69 +44,20 @@ class TorrentService:
 
     def download(self, indexer_result: IndexerQueryResult) -> Torrent:
         log.info(f"Attempting to download torrent: {indexer_result.title}")
-        torrent = Torrent(
-            status=TorrentStatus.unknown,
-            title=indexer_result.title,
-            quality=indexer_result.quality,
-            imported=False,
-            hash="",
-        )
 
-        url = indexer_result.download_url
-        torrent_filepath = BasicConfig().torrent_directory / f"{torrent.title}.torrent"
+        # Use the download client to handle the download and get the complete torrent object
+        torrent = self.download_client.download_torrent(indexer_result)
 
-        if torrent_filepath.exists():
-            log.warning(f"Torrent already exists: {torrent_filepath}")
-            return self.get_torrent_status(torrent=torrent)
-
-        with open(torrent_filepath, "wb") as file:
-            content = requests.get(url).content
-            file.write(content)
-
-        with open(torrent_filepath, "rb") as file:
-            content = file.read()
-            try:
-                decoded_content = bencoder.decode(content)
-            except Exception as e:
-                log.error(f"Failed to decode torrent file: {e}")
-                raise e
-            torrent.hash = hashlib.sha1(
-                bencoder.encode(decoded_content[b"info"])
-            ).hexdigest()
-            answer = self.api_client.torrents_add(
-                category="MediaManager", torrent_files=content, save_path=torrent.title
-            )
-
-        if answer == "Ok.":
-            log.info(f"Successfully added torrent: {torrent.title}")
-            return self.get_torrent_status(torrent=torrent)
-        else:
-            log.error(f"Failed to download torrent. API response: {answer}")
-            raise RuntimeError(
-                f"Failed to download torrent, API-Answer isn't 'Ok.'; API Answer: {answer}"
-            )
+        # Save to repository and return
+        return self.torrent_repository.save_torrent(torrent=torrent)
 
     def get_torrent_status(self, torrent: Torrent) -> Torrent:
         log.info(f"Fetching status for torrent: {torrent.title}")
-        info = self.api_client.torrents_info(torrent_hashes=torrent.hash)
 
-        if not info:
-            log.warning(f"No information found for torrent: {torrent.id}")
-            torrent.status = TorrentStatus.unknown
-        else:
-            state: str = info[0]["state"]
-            log.info(f"Torrent {torrent.id} is in state: {state}")
+        # Get status from download client
+        torrent.status = self.download_client.get_torrent_status(torrent)
 
-            if state in self.DOWNLOADING_STATE:
-                torrent.status = TorrentStatus.downloading
-            elif state in self.FINISHED_STATE:
-                torrent.status = TorrentStatus.finished
-            elif state in self.ERROR_STATE:
-                torrent.status = TorrentStatus.error
-            elif state in self.UNKNOWN_STATE:
-                torrent.status = TorrentStatus.unknown
-            else:
-                torrent.status = TorrentStatus.error
+        # Save updated status to repository
         self.torrent_repository.save_torrent(torrent=torrent)
         return torrent
 
@@ -161,7 +69,7 @@ class TorrentService:
         :param torrent: the torrent to cancel
         """
         log.info(f"Cancelling download for torrent: {torrent.title}")
-        self.api_client.torrents_delete(delete_files=delete_files)
+        self.download_client.remove_torrent(torrent, delete_data=delete_files)
         return self.get_torrent_status(torrent=torrent)
 
     def pause_download(self, torrent: Torrent) -> Torrent:
@@ -171,7 +79,7 @@ class TorrentService:
         :param torrent: the torrent to pause
         """
         log.info(f"Pausing download for torrent: {torrent.title}")
-        self.api_client.torrents_pause(torrent_hashes=torrent.hash)
+        self.download_client.pause_torrent(torrent)
         return self.get_torrent_status(torrent=torrent)
 
     def resume_download(self, torrent: Torrent) -> Torrent:
@@ -181,7 +89,7 @@ class TorrentService:
         :param torrent: the torrent to resume
         """
         log.info(f"Resuming download for torrent: {torrent.title}")
-        self.api_client.torrents_resume(torrent_hashes=torrent.hash)
+        self.download_client.resume_torrent(torrent)
         return self.get_torrent_status(torrent=torrent)
 
     def get_all_torrents(self) -> list[Torrent]:
@@ -202,5 +110,6 @@ class TorrentService:
     #        from media_manager.tv.repository import remove_season_files_by_torrent_id
     #        remove_season_files_by_torrent_id(db=self.db, torrent_id=torrent_id)
     #    media_manager.torrent.repository.delete_torrent(db=self.db, torrent_id=t.id)
+
     def get_movie_files_of_torrent(self, torrent: Torrent):
         return self.torrent_repository.get_movie_files_of_torrent(torrent_id=torrent.id)