add download controls to deluge client

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

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
 

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