def test_compute_bucket_number(self):
'''It should return the bucket number based on leading zeros bits'''
self.assertEqual(compute_bucket_number(
KeyBytes('00' * 20), KeyBytes('00' * 20)), 160)
self.assertEqual(compute_bucket_number(
KeyBytes('00' * 20), KeyBytes('FF' * 20)), 0)
self.assertEqual(compute_bucket_number(
KeyBytes('00' * 20), KeyBytes('00' * 10 + 'FF' * 10)), 80)
def get_bucket_number(self, node):
'''Get the appropriate bucket number for the node
:rtype: `int`
'''
return compute_bucket_number(self._key, node.key)
def test_random_bucket_key(self):
'''It should generate keys that goes into given bucket number'''
node_key = KeyBytes()
for i in range(160):
key = random_bucket_key(node_key, i)
self.assertEqual(i, compute_bucket_number(node_key, key))
def count_close(self, key):
'''Return the number of node closer than the given key'''
bucket = self._buckets[compute_bucket_number(self._key, key)]
count = 0
for node in bucket:
if node.key.distance_int(key) < self._key.distance_int(key):
count += 1
return count
def _calculate_expiration_time(self, key):
'''Return the expiration time for a given key'''
bucket_number = compute_bucket_number(self.key, key)
num_contacts = sum(
[len(self.routing_table[i]) for i in range(bucket_number)])
num_bucket_contacts = self._routing_table.count_close(key)
c = num_contacts + num_bucket_contacts
if c < Bucket.MAX_BUCKET_SIZE == 0:
return DHTNetwork.TIME_EXPIRE
else:
return DHTNetwork.TIME_EXPIRE / math.exp(
c / Bucket.MAX_BUCKET_SIZE)
def get_close_nodes(self, key, count=3):
'''Return the closest nodes to a key
:Parameters:
key : :class:`.KeyBytes`
The target key
count: `int`
The maximum length of the list returned
:return: A ``list`` of `Node`
:rtype: ``list``
'''
bucket_number = compute_bucket_number(self._key, key)
bucket = self._buckets[bucket_number]
if len(bucket) >= count:
return random.sample(bucket.nodes, count)
# Pick nodes from random buckets
nodes = set(bucket.nodes)
buckets = list(self._buckets)
random.shuffle(buckets)
for bucket in buckets:
num_needed = min(len(bucket), count - len(nodes))
for contact in random.sample(bucket.nodes, num_needed):
nodes.add(contact)
if len(nodes) == count:
break
assert len(nodes) <= count
_logger.debug('Got %s close nodes', len(nodes))
return list(nodes)
