refactor: introduce TorrentInfo dataclass and enhance testing

feral_services/jackett.py

@@ -2,14 +2,35 @@
 
 import os
 import random
+from dataclasses import dataclass
 
 import requests
 from requests.exceptions import Timeout
 
-
 _TOTAL_RESULTS_TO_RETURN = 20
 
 
+@dataclass
+class TorrentInfo:
+    name: str
+    size: str
+    seeds: int
+    peers: int
+    source: str
+    magnet: str
+    link: str
+
+    def format_response(self, req_id: str | None = None) -> str:
+        prefix = f"/get{req_id} - " if req_id else "Success - "
+        return (
+            f"{prefix}{self.name}\n"
+            f"└─ {self.source} | "
+            f"Seeds: {self.seeds:,} | "
+            f"Peers: {self.peers:,} | "
+            f"Size: {self.size}"
+        )
+
+
 def search(query) -> (str, list):
     params = (
         ('apikey', os.getenv('JACKETT_API_KEY')),
@@ -67,28 +88,19 @@ def format_and_filter_results(results: list, user_id: int, user_id_to_results: d
             if count > 4:
                 return 'id collision happened?...'
 
-        user_id_to_results[user_id][req_id] = {
-            'magnet': result['MagnetUri'],
-            'link': result['Link'],
-            'label': result['Tracker'],
-            'title': result['Title'],
-            'size': round((result['Size'] / 1024 / 1024 / 1024), 2),
-        }
-
-        details = (
-            f"└─ {result['Tracker']} | "
-            f"Seeds: {format(result['Seeders'], ',')} | "
-            f"Peers: {format(result['Peers'], ',')} | "
-            f"Size: {format(round(result['Size'] / 1024 / 1024 / 1024, 2), '.2f')} GB"
-        )
-
-        formatted_result = (
-            f"/get{req_id} - {result['Title']}\n"
-            f"{details}"
+        torrent_info = TorrentInfo(
+            name=result['Title'],
+            size=f"{round((result['Size'] / 1024 / 1024 / 1024), 2)} GB",
+            seeds=result['Seeders'],
+            peers=result['Peers'],
+            source=result['Tracker'],
+            magnet=result['MagnetUri'],
+            link=result['Link']
         )
 
-        returned_results.append(formatted_result)
+        user_id_to_results[user_id][req_id] = torrent_info
+        returned_results.append(torrent_info.format_response(req_id))
 
-    result_count = f"Results ({len(returned_results)}/{len(results)})"
+    result_count_str = f'Results ({len(returned_results)}/{len(results)})'
     returned_results_str = '\n\n'.join(returned_results)
-    return f"{result_count}\n\n{returned_results_str}"
+    return f'{result_count_str}\n\n{returned_results_str}'

feral_services/ru_torrent.py

@@ -7,6 +7,8 @@
 import bencodepy
 import requests
 
+from feral_services.jackett import TorrentInfo
+
 load_dotenv()
 
 
@@ -17,10 +19,10 @@
 def _format_return_url(url):
     parsed_url = urllib.parse.urlparse(url)
     query_params = urllib.parse.parse_qs(parsed_url.query)
-    return query_params['result[]'][0]
+    return query_params['result[]'][0].strip().casefold()
 
 
-def upload_torrent(torrent_file, label, username):
+def upload_torrent(torrent_file, label: str, username: str, torrent_info: TorrentInfo):
     metadata = bencodepy.decode(torrent_file)
     file_name = urllib.parse.quote(
         metadata[b'info'][b'name'].decode(),  # noqa
@@ -37,11 +39,13 @@ def upload_torrent(torrent_file, label, username):
     )
 
     if response.ok:
-        return _format_return_url(response.url)
+        status = _format_return_url(response.url)
+        if status == "success":
+            return torrent_info.format_response()
     return f'Error: {response.status_code}'
 
 
-def upload_magnet(magnet_link, label, username):
+def upload_magnet(magnet_link: str, label: str, username: str, torrent_info: TorrentInfo | None = None):
     if username:
         label = f'{username}, {label}'
 
@@ -53,5 +57,10 @@ def upload_magnet(magnet_link, label, username):
     )
 
     if response.ok:
-        return _format_return_url(response.url)
-    return f'Error: {response.status_code}'
+        status = _format_return_url(response.url)
+        if torrent_info:
+            if status == "success":
+                return torrent_info.format_response()
+        else:
+            return status.capitalize()
+    return f'Error: {response.status_code}'

main.py

@@ -16,6 +16,8 @@
 from feral_services import jackett
 from feral_services import ru_torrent
 from feral_services.instance import execute_command
+from feral_services.jackett import TorrentInfo
+
 load_dotenv()
 
 _RESULTS = {}
@@ -147,22 +149,23 @@ async def get(update: Update, _context):
         return
 
     _, get_id = update.message.text.split('/get', 1)
-    result = users_data.get(get_id)
+    result: TorrentInfo = users_data.get(get_id)
     if not result:
         await update.message.reply_text('Not a valid item')
         return
 
     username = update.effective_user.username or update.effective_user.first_name
-    if magnet := result.get('magnet'):
+    if magnet := result.magnet:
         magnet_upload_result = ru_torrent.upload_magnet(
             magnet,
-            result['label'],
+            result.source,
             username,
+            result
         )
         await update.message.reply_text(magnet_upload_result)
         return
 
-    elif link := result.get('link'):
+    elif link := result.link:
         url_response = requests.get(link, allow_redirects=False)
         if not url_response.ok:
             await update.message.reply_text(
@@ -171,23 +174,28 @@ async def get(update: Update, _context):
             )
             return
 
-        with contextlib.suppress(Exception):
+        try:
             if url_response.status_code == 302:
                 magnet_upload_result = ru_torrent.upload_magnet(
                     url_response.headers['Location'],
-                    result['label'],
+                    result.source,
                     username,
+                    result
                 )
                 await update.message.reply_text(magnet_upload_result)
                 return
 
             torrent_upload_result = ru_torrent.upload_torrent(
                 url_response.content,
-                result['label'],
+                result.source,
                 username,
+                result
             )
             await update.message.reply_text(torrent_upload_result)
             return
+        except Exception as e:
+            print(e)
+
 
     await update.message.reply_text('Something went wrong')
 

tests/jackett_test.py

@@ -87,7 +87,28 @@ def test_format_and_filter_results(mocker):
 
     user_id = 12345
     memory_database = {}
-    results = _SHAWSHANK_RESULTS
+
+    results = [
+        {
+            'Title': 'The Shawshank Redemption 1994 REMASTERED 1080p BluRay H264 AAC-LAMA',
+            'Size': 2910237066,  # ~2.71GB
+            'Seeders': 42,
+            'Peers': 0,
+            'Tracker': 'IPTorrents',
+            'MagnetUri': 'magnet:test1',
+            'Link': 'http://test1.com'
+        },
+        {
+            'Title': 'The Shawshank Redemption 1994 REMASTERED 1080p BluRay H264 AAC R4RBG TGx',
+            'Size': 2984725094,  # ~2.78GB
+            'Seeders': 84,
+            'Peers': 15,
+            'Tracker': '1337x',
+            'MagnetUri': 'magnet:test2',
+            'Link': 'http://test2.com'
+        }
+    ]
+
     formatted_results = jackett.format_and_filter_results(
         results, user_id, memory_database,
     )
@@ -96,20 +117,23 @@ def test_format_and_filter_results(mocker):
 
     expected_result_count_str = 'Results (2/2)'
     expected_returned_results_str = (
-        '/get11111 - IPTorrents, Seeds: 42, Peers: 0, Size: 2.71GB\n'
-        'The Shawshank Redemption 1994 REMASTERED 1080p BluRay H264 AAC-LAMA\n\n'  # noqa
-        '/get22222 - 1337x, Seeds: 84, Peers: 15, Size: 2.78GB\n'
-        'The Shawshank Redemption 1994 REMASTERED 1080p BluRay H264 AAC R4RBG TGx'  # noqa
+        '/get11111 - The Shawshank Redemption 1994 REMASTERED 1080p BluRay H264 AAC-LAMA\n'
+        '└─ IPTorrents | Seeds: 42 | Peers: 0 | Size: 2.71 GB\n\n'
+        '/get22222 - The Shawshank Redemption 1994 REMASTERED 1080p BluRay H264 AAC R4RBG TGx\n'
+        '└─ 1337x | Seeds: 84 | Peers: 15 | Size: 2.78 GB'
     )
+
     assert formatted_results.startswith(expected_result_count_str)
     assert expected_returned_results_str in formatted_results
 
     assert len(memory_database) == 1
     assert isinstance(memory_database[user_id], dict)
     assert len(memory_database[user_id]) == 2
-    for req_id, values in memory_database[user_id].items():
+
+    for req_id, torrent_info in memory_database[user_id].items():
         assert isinstance(req_id, str)
-        assert isinstance(values, dict)
-        assert set(values.keys()) == {
-            'magnet', 'link', 'label', 'title', 'size',
-        }
+        assert isinstance(torrent_info, jackett.TorrentInfo)
+        assert all(
+            hasattr(torrent_info, attr)
+            for attr in ['name', 'size', 'seeds', 'peers', 'source', 'magnet', 'link']
+        )