diff --git a/feral_services/jackett.py b/feral_services/jackett.py index eeec0e9..4681201 100644 --- a/feral_services/jackett.py +++ b/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}' \ No newline at end of file diff --git a/feral_services/ru_torrent.py b/feral_services/ru_torrent.py index 7b15a81..0b9bff8 100644 --- a/feral_services/ru_torrent.py +++ b/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}' \ No newline at end of file diff --git a/main.py b/main.py index 2df5fe2..0cef389 100644 --- a/main.py +++ b/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') diff --git a/tests/jackett_test.py b/tests/jackett_test.py index 485b845..5e809c6 100644 --- a/tests/jackett_test.py +++ b/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'] + ) \ No newline at end of file diff --git a/tests/ru_torrent_test.py b/tests/ru_torrent_test.py index f5a9156..a67695c 100644 --- a/tests/ru_torrent_test.py +++ b/tests/ru_torrent_test.py @@ -11,17 +11,27 @@ @pytest.mark.parametrize( 'url, expected_result', [ - ('https://example.com?result[]=Success', 'Success'), - ('https://example.com?result[]=Success&result[]=baz', 'Success'), + ('https://example.com?result[]=Success', 'success'), + ('https://example.com?result[]=Success&result[]=baz', 'success'), + ('https://example.com?result[]=FAILURE', 'failure'), + ('https://example.com?result[]=ERROR', 'error'), + ('https://example.com?result[]= SUCCESS ', 'success'), # Test whitespace handling ], ) def test_format_return_url(url, expected_result): assert ru_torrent._format_return_url(url) == expected_result -def test_format_return_url_key_error(): +@pytest.mark.parametrize( + 'url', [ + 'https://example.com', + 'https://example.com?other[]=Success', + 'https://example.com?result=Success', # Missing brackets + ] +) +def test_format_return_url_key_error(url): with pytest.raises(KeyError): - ru_torrent._format_return_url('https://example.com') + ru_torrent._format_return_url(url) class MockResponse: @@ -31,28 +41,141 @@ def __init__(self, status_code, ok=True, url=None): self.url = url +class MockTorrentInfo: + def format_response(self): + return "formatted_response" + + +@pytest.fixture +def mock_env(mocker): + mocker.patch.dict( + 'os.environ', { + 'RU_TORRENT_TOKEN': 'test_token', + 'RU_TORRENT_URL': 'http://example.com' + } + ) + + +@pytest.fixture +def mock_torrent_info(): + return MockTorrentInfo() + + @pytest.mark.parametrize( - 'upload_function', [ - ru_torrent.upload_torrent, - ru_torrent.upload_magnet, - ], + 'username, label, expected_label', [ + ('test_user', 'test_label', 'test_user, test_label'), + ('', 'test_label', 'test_label'), + (None, 'test_label', 'test_label'), + ('test_user', '', 'test_user, '), + ('test_user', 'label with spaces', 'test_user, label with spaces'), + ] ) +def test_upload_torrent_label_formatting(mocker, mock_env, mock_torrent_info, username, label, expected_label): + mock_response = MockResponse(status_code=200, ok=True, url='http://example.com?result[]=Success') + mocker.patch.object(requests, 'post', return_value=mock_response) + mocker.patch.object(bencodepy, 'decode', return_value={b'info': {b'name': b'test.torrent'}}) + mocker.patch.object(urllib.parse, 'quote', return_value='test.torrent') + + ru_torrent.upload_torrent('fake_torrent', label, username, mock_torrent_info) + + _, kwargs = requests.post.call_args + assert kwargs['params']['label'] == expected_label + + @pytest.mark.parametrize( - 'response', [ - MockResponse( - status_code=200, - url='http://example.com?result[]=Success', - ), - MockResponse(status_code=400, ok=False), - ], + 'torrent_name, quoted_name', [ + (b'test.torrent', 'test.torrent'), + (b'test with spaces.torrent', 'test%20with%20spaces.torrent'), + (b'test+special&chars.torrent', 'test%2Bspecial%26chars.torrent'), + ] ) -def test_upload_function(mocker, upload_function, response): - mocker.patch.object(requests, 'post', return_value=response) - mocker.patch.object(bencodepy, 'decode') - mocker.patch.object(urllib.parse, 'quote') - - result = upload_function('file_or_link', 'my_label', 'my_username') - if response.ok: - assert result == ru_torrent._format_return_url(response.url) +def test_upload_torrent_name_encoding(mocker, mock_env, mock_torrent_info, torrent_name, quoted_name): + mock_response = MockResponse(status_code=200, ok=True, url='http://example.com?result[]=Success') + mocker.patch.object(requests, 'post', return_value=mock_response) + mocker.patch.object(bencodepy, 'decode', return_value={b'info': {b'name': torrent_name}}) + mocker.patch.object(urllib.parse, 'quote', side_effect=urllib.parse.quote) + + ru_torrent.upload_torrent('fake_torrent', 'label', 'user', mock_torrent_info) + + _, kwargs = requests.post.call_args + assert kwargs['files']['torrent_file'][0] == quoted_name + + +@pytest.mark.parametrize( + 'response_data', [ + {'status_code': 200, 'ok': True, 'url': 'http://example.com?result[]=Success'}, + {'status_code': 400, 'ok': False}, + {'status_code': 500, 'ok': False}, + {'status_code': 200, 'ok': True, 'url': 'http://example.com?result[]=Failure'}, + ] +) +def test_upload_torrent_responses(mocker, mock_env, mock_torrent_info, response_data): + mock_response = MockResponse(**response_data) + mocker.patch.object(requests, 'post', return_value=mock_response) + mocker.patch.object(bencodepy, 'decode', return_value={b'info': {b'name': b'test.torrent'}}) + mocker.patch.object(urllib.parse, 'quote', return_value='test.torrent') + + result = ru_torrent.upload_torrent('fake_torrent', 'label', 'user', mock_torrent_info) + + if not mock_response.ok: + assert result == f'Error: {mock_response.status_code}' + elif 'url' in response_data and 'Success' in response_data['url']: + assert result == "formatted_response" else: - assert result == f'Error: {response.status_code}' + assert 'Error' in result + + +@pytest.mark.parametrize( + 'username, label, expected_label', [ + ('test_user', 'test_label', 'test_user, test_label'), + ('', 'test_label', 'test_label'), + (None, 'test_label', 'test_label'), + ('test_user', '', 'test_user, '), + ('test_user', 'label with spaces', 'test_user, label with spaces'), + ] +) +def test_upload_magnet_label_formatting(mocker, mock_env, username, label, expected_label): + mock_response = MockResponse(status_code=200, ok=True, url='http://example.com?result[]=Success') + mocker.patch.object(requests, 'post', return_value=mock_response) + + ru_torrent.upload_magnet('magnet:?xt=test', label, username) + + _, kwargs = requests.post.call_args + assert kwargs['params']['label'] == expected_label + + +@pytest.mark.parametrize( + 'response_data, torrent_info, expected_result', [ + ( + {'status_code': 200, 'ok': True, 'url': 'http://example.com?result[]=Success'}, + MockTorrentInfo(), + "formatted_response" + ), + ( + {'status_code': 200, 'ok': True, 'url': 'http://example.com?result[]=Success'}, + None, + "Success" + ), + ( + {'status_code': 400, 'ok': False}, + MockTorrentInfo(), + "Error: 400" + ), + ( + {'status_code': 400, 'ok': False}, + None, + "Error: 400" + ), + ( + {'status_code': 200, 'ok': True, 'url': 'http://example.com?result[]=Failure'}, + MockTorrentInfo(), + "Error: 200" + ), + ] +) +def test_upload_magnet_responses(mocker, mock_env, response_data, torrent_info, expected_result): + mock_response = MockResponse(**response_data) + mocker.patch.object(requests, 'post', return_value=mock_response) + + result = ru_torrent.upload_magnet('magnet:?xt=test', 'label', 'user', torrent_info) + assert result == expected_result \ No newline at end of file