Add binary search to get bucket for node

evm/p2p/kademlia.py

@@ -191,10 +191,13 @@ def __contains__(self, node):
     def __len__(self):
         return len(self.nodes)
 
-    def __le__(self, other):
+    def __lt__(self, other):
         if not isinstance(other, self.__class__):
             raise TypeError("Cannot compare KBucket with type {}.".format(other.__class__))
-        return self.end <= other.end
+        # Check for invalid state of KBuckets
+        if not self.end < other.start:
+            raise ValueError("Invalid Buckets.")
+        return self.end < other.start
 
 
 class RoutingTable:
@@ -270,15 +273,14 @@ def neighbours(self, node_id, k=k_bucket_size):
 
 def binary_get_bucket_for_node(buckets, node):
     """Return the bucket for a given node."""
-    sorted_buckets = sorted(buckets)
-    bucket_ends = [bucket.end for bucket in sorted_buckets]
+    bucket_ends = [bucket.end for bucket in buckets]
     bucket_position = bisect.bisect_left(bucket_ends, node.id)
     # Prevents edge cases where bisect_left returns an out of range index
     try:
-        bucket = sorted_buckets[bucket_position]
+        bucket = buckets[bucket_position]
     except IndexError:
         raise ValueError("No bucket found for node with id {}".format(node.id))
-    bucket = sorted_buckets[bucket_position]
+    bucket = buckets[bucket_position]
     if not bucket.start <= node.id <= bucket.end:
         raise ValueError("No bucket found for node with id {}".format(node.id))
     return bucket

evm/p2p/test_kademlia.py

@@ -264,52 +264,56 @@ def test_kbucket_split():
 
 
 def test_bucket_ordering():
-    first = kademlia.KBucket(51, 100)
-    second = kademlia.KBucket(0, 50)
+    first = kademlia.KBucket(0, 50)
+    second = kademlia.KBucket(51, 100)
     third = random_node()
-    assert first > second
+    assert first < second
     with pytest.raises(TypeError):
             first > third
 
 
 @pytest.mark.parametrize(
-    "bucket_list, node_id, correct",
+    "bucket_list, node_id",
     (
-        (list([]), 5, None),
+        (list([]), 5),
         # test for node.id < bucket.end
-        (list([kademlia.KBucket(0, 4)]), 5, None),
+        (list([kademlia.KBucket(0, 4)]), 5),
         # test for node.id > bucket.start
-        (list([kademlia.KBucket(6, 10)]), 5, None),
+        (list([kademlia.KBucket(6, 10)]), 5),
         # test multiple buckets that don't contain node.id
         (list(
             [
-                kademlia.KBucket(0, 1),
+                kademlia.KBucket(1, 5),
+                kademlia.KBucket(6, 49),
                 kademlia.KBucket(50, 100),
-                kademlia.KBucket(6, 49)
             ]
-        ), 5, None),
-        # test single bucket
+        ), 0),
+    )
+)
+def test_binary_get_bucket_for_node_error(bucket_list, node_id):
+    node = random_node(nodeid=node_id)
+    with pytest.raises(ValueError):
+        kademlia.binary_get_bucket_for_node(bucket_list, node)
+
+
+@pytest.mark.parametrize(
+    "bucket_list, node_id, correct_position",
+    (
         (list([kademlia.KBucket(0, 100)]), 5, 0),
-        (list([kademlia.KBucket(50, 100), kademlia.KBucket(0, 49)]), 5, 1),
-        # test multiple unordered buckets
+        (list([kademlia.KBucket(0, 49), kademlia.KBucket(50, 100)]), 5, 0),
         (list(
             [
-                kademlia.KBucket(50, 100),
                 kademlia.KBucket(0, 1),
                 kademlia.KBucket(2, 5),
-                kademlia.KBucket(6, 49)
+                kademlia.KBucket(6, 49),
+                kademlia.KBucket(50, 100)
             ]
-        ), 5, 2),
+        ), 5, 1),
     )
 )
-def test_binary_get_bucket_for_node(bucket_list, node_id, correct):
-    node = random_node()
-    node.id = node_id
-    if correct is None:
-        with pytest.raises(ValueError):
-            kademlia.binary_get_bucket_for_node(bucket_list, node)
-    else:
-        assert kademlia.binary_get_bucket_for_node(bucket_list, node) == bucket_list[correct]
+def test_binary_get_bucket_for_node(bucket_list, node_id, correct_position):
+    node = random_node(nodeid=node_id)
+    assert kademlia.binary_get_bucket_for_node(bucket_list, node) == bucket_list[correct_position]
 
 
 def test_compute_shared_prefix_bits():