From a5127da5f20ff6dbe76701de1c081dab937f7371 Mon Sep 17 00:00:00 2001
From: gmega <giuliano.mega@gmail.com>
Date: Sun, 3 Nov 2024 08:39:19 -0300
Subject: [PATCH] add download synchronization to base static experiment

---
 benchmarks/experiments/static_experiment.py   | 20 +++++++++++++------
 .../tests/test_static_experiment.py           | 16 +++++++++++++--
 2 files changed, 28 insertions(+), 8 deletions(-)

diff --git a/benchmarks/experiments/static_experiment.py b/benchmarks/experiments/static_experiment.py
index 906c1d7..11062a2 100644
--- a/benchmarks/experiments/static_experiment.py
+++ b/benchmarks/experiments/static_experiment.py
@@ -19,17 +19,25 @@ class StaticDisseminationExperiment(Generic[TNetworkHandle, TInitialMetadata]):
 
     def run(self):
         sample = self.sampler(len(self.network.nodes))
-        seeder_idx = [next(sample) for _ in range(0, self.seeders)]
+        seeder_indexes = [next(sample) for _ in range(0, self.seeders)]
         seeders, leechers = (
-            [self.network.nodes[i] for i in seeder_idx],
-            [self.network.nodes[i] for i in range(0, len(self.network.nodes)) if i not in seeder_idx]
+            [
+                self.network.nodes[i]
+                for i in seeder_indexes
+            ],
+            [
+                self.network.nodes[i]
+                for i in range(0, len(self.network.nodes))
+                if i not in seeder_indexes
+            ]
         )
 
         meta, data = self.generate_data()
-        handle = meta
 
+        handle = meta
         for node in seeders:
             handle = node.seed(data, handle)
 
-        for node in leechers:
-            node.leech(handle)
+        handles = [node.leech(handle) for node in leechers]
+        for handle in handles:
+            handle.await_for_completion()
diff --git a/benchmarks/experiments/tests/test_static_experiment.py b/benchmarks/experiments/tests/test_static_experiment.py
index 1878323..a1495ca 100644
--- a/benchmarks/experiments/tests/test_static_experiment.py
+++ b/benchmarks/experiments/tests/test_static_experiment.py
@@ -2,7 +2,7 @@ from dataclasses import dataclass
 from pathlib import Path
 from typing import Optional, List, Tuple, Union, Sequence
 
-from benchmarks.core.network import FileSharingNetwork, Node
+from benchmarks.core.network import FileSharingNetwork, Node, DownloadHandle
 from benchmarks.core.utils import Sampler
 from benchmarks.experiments.static_experiment import StaticDisseminationExperiment
 
@@ -22,6 +22,7 @@ class MockNode(Node[MockHandle, str]):
     def __init__(self) -> None:
         self.seeding: Optional[Tuple[MockHandle, Path]] = None
         self.leeching: Optional[MockHandle] = None
+        self.download_was_awaited = False
 
     def seed(
             self,
@@ -37,6 +38,16 @@ class MockNode(Node[MockHandle, str]):
 
     def leech(self, handle: MockHandle):
         self.leeching = handle
+        return MockDownloadHandle(self)
+
+
+class MockDownloadHandle(DownloadHandle):
+    def __init__(self, parent: MockNode) -> None:
+        self.parent = parent
+
+    def await_for_completion(self, timeout: float = 0) -> bool:
+        self.parent.download_was_awaited = True
+        return True
 
 
 class MockFileSharingNetwork(FileSharingNetwork[MockHandle, str]):
@@ -72,7 +83,7 @@ def test_should_place_seeders():
     assert actual_seeders == set(seeder_indexes)
 
 
-def test_should_place_leechers():
+def test_should_download_at_remaining_nodes():
     network = MockFileSharingNetwork(n=13)
     file = Path('/path/to/data')
     seeder_indexes = [9, 6, 3]
@@ -92,6 +103,7 @@ def test_should_place_leechers():
             assert node.leeching.path == file
             assert node.leeching.name == 'data'
             assert node.seeding is None
+            assert node.download_was_awaited
             actual_leechers.add(index)
 
     assert actual_leechers == set(range(13)) - set(seeder_indexes)