diff --git a/benchmarks/core/deluge.py b/benchmarks/core/deluge.py index b69830f..b2b03e4 100644 --- a/benchmarks/core/deluge.py +++ b/benchmarks/core/deluge.py @@ -1,8 +1,10 @@ import base64 +import logging import shutil from dataclasses import dataclass from io import BytesIO from pathlib import Path +from time import time from typing import List, Union, Optional, Self, Dict, Any import pathvalidate @@ -10,7 +12,9 @@ from deluge_client import DelugeRPCClient from torrentool.torrent import Torrent from urllib3.util import Url -from benchmarks.core.network import TNetworkHandle, SharedFSNode +from benchmarks.core.network import SharedFSNode, DownloadHandle + +logger = logging.getLogger(__name__) @dataclass(frozen=True) @@ -85,8 +89,17 @@ class DelugeNode(SharedFSNode[Torrent, DelugeMeta]): return torrent - def leech(self, handle: TNetworkHandle): - pass + def leech(self, handle: Torrent) -> DownloadHandle: + self.rpc.core.add_torrent_file( + filename=f'{handle.name}.torrent', + filedump=self._b64dump(handle), + options=dict(), + ) + + return DelugeDownloadHandle( + node=self, + torrent=handle, + ) def torrent_info(self, name: str) -> List[Dict[bytes, Any]]: return list(self.rpc.core.get_torrents_status({'name': name}, []).values()) @@ -111,3 +124,24 @@ class DelugeNode(SharedFSNode[Torrent, DelugeMeta]): buffer = BytesIO() buffer.write(handle.to_string()) return base64.b64encode(buffer.getvalue()) + + +class DelugeDownloadHandle(DownloadHandle): + + def __init__(self, torrent: Torrent, node: DelugeNode) -> None: + self.node = node + self.torrent = torrent + + def await_for_completion(self, timeout: float = 0) -> bool: + name = self.torrent.name + current = time() + while (time() - current) <= timeout: + response = self.node.rpc.core.get_torrents_status({'name': name}, []) + if len(response) > 1: + logger.warning(f'Client has multiple torrents matching name {name}. Returning the first one.') + + status = list(response.values())[0] + if status[b'is_finished']: + return True + + return False diff --git a/benchmarks/core/network.py b/benchmarks/core/network.py index 8e1a91f..c3b1acd 100644 --- a/benchmarks/core/network.py +++ b/benchmarks/core/network.py @@ -9,6 +9,18 @@ TNetworkHandle = TypeVar('TNetworkHandle') TInitialMetadata = TypeVar('TInitialMetadata') +class DownloadHandle(ABC): + """A :class:`DownloadHandle` represents a reference to an underlying download.""" + + @abstractmethod + def await_for_completion(self, timeout: float = 0) -> bool: + """Blocks the current thread until either the download completes or a timeout expires. + + :param timeout: Timeout in seconds. + :return: True if the download completed within the timeout, False otherwise.""" + pass + + class Node(ABC, Generic[TNetworkHandle, TInitialMetadata]): """A :class:`Node` represents a peer within a :class:`FileSharingNetwork`.""" @@ -31,7 +43,7 @@ class Node(ABC, Generic[TNetworkHandle, TInitialMetadata]): pass @abstractmethod - def leech(self, handle: TNetworkHandle): + def leech(self, handle: TNetworkHandle) -> DownloadHandle: """Makes the current node a leecher for the provided handle.""" pass diff --git a/benchmarks/core/tests/test_deluge_node.py b/benchmarks/core/tests/test_deluge_node.py index 2c75e7a..4633b3f 100644 --- a/benchmarks/core/tests/test_deluge_node.py +++ b/benchmarks/core/tests/test_deluge_node.py @@ -9,9 +9,8 @@ from benchmarks.core.utils import megabytes from benchmarks.tests.utils import shared_volume -@pytest.fixture -def deluge_node1() -> Generator[DelugeNode, None, None]: - node = DelugeNode('deluge-1', volume=shared_volume(), daemon_port=6890) +def deluge_node(name: str, port: int) -> Generator[DelugeNode, None, None]: + node = DelugeNode(name, volume=shared_volume(), daemon_port=port) node.wipe_all_torrents() try: yield node @@ -19,6 +18,16 @@ def deluge_node1() -> Generator[DelugeNode, None, None]: node.wipe_all_torrents() +@pytest.fixture +def deluge_node1() -> Generator[DelugeNode, None, None]: + yield from deluge_node('deluge-1', 6890) + + +@pytest.fixture +def deluge_node2() -> Generator[DelugeNode, None, None]: + yield from deluge_node('deluge-2', 6893) + + def test_should_seed_files(deluge_node1: DelugeNode, temp_random_file: Path, tracker: Url): assert not deluge_node1.torrent_info(name='dataset1') @@ -30,3 +39,24 @@ def test_should_seed_files(deluge_node1: DelugeNode, temp_random_file: Path, tra assert info[b'name'] == b'dataset1' assert info[b'total_size'] == megabytes(1) assert info[b'is_seed'] == True + + +def test_should_download_files( + deluge_node1: DelugeNode, deluge_node2: DelugeNode, + temp_random_file: Path, tracker: Url): + + assert not deluge_node1.torrent_info(name='dataset1') + assert not deluge_node2.torrent_info(name='dataset1') + + torrent = deluge_node1.seed(temp_random_file, DelugeMeta(name='dataset1', announce_url=tracker)) + handle = deluge_node2.leech(torrent) + + assert handle.await_for_completion(5) + + response = deluge_node2.torrent_info(name='dataset1') + assert len(response) == 1 + info = response[0] + + assert info[b'name'] == b'dataset1' + assert info[b'total_size'] == megabytes(1) + assert info[b'is_seed'] == True \ No newline at end of file