forked from LBRYCommunity/lbry-sdk
143 lines
6.2 KiB
Python
143 lines
6.2 KiB
Python
import struct
|
|
import asyncio
|
|
from lbrynet.utils import generate_id
|
|
from lbrynet.dht.protocol.routing_table import KBucket
|
|
from lbrynet.dht.peer import PeerManager, KademliaPeer
|
|
from lbrynet.dht import constants
|
|
from torba.testcase import AsyncioTestCase
|
|
|
|
|
|
def address_generator(address=(10, 42, 42, 1)):
|
|
def increment(addr):
|
|
value = struct.unpack("I", "".join([chr(x) for x in list(addr)[::-1]]).encode())[0] + 1
|
|
new_addr = []
|
|
for i in range(4):
|
|
new_addr.append(value % 256)
|
|
value >>= 8
|
|
return tuple(new_addr[::-1])
|
|
|
|
while True:
|
|
yield "{}.{}.{}.{}".format(*address)
|
|
address = increment(address)
|
|
|
|
|
|
class TestKBucket(AsyncioTestCase):
|
|
def setUp(self):
|
|
self.loop = asyncio.get_event_loop()
|
|
self.address_generator = address_generator()
|
|
self.peer_manager = PeerManager(self.loop)
|
|
self.kbucket = KBucket(self.peer_manager, 0, 2**constants.hash_bits, generate_id())
|
|
|
|
def test_add_peer(self):
|
|
peer = KademliaPeer(None, '1.2.3.4', constants.generate_id(2), udp_port=4444)
|
|
peer_update2 = KademliaPeer(None, '1.2.3.4', constants.generate_id(2), udp_port=4445)
|
|
|
|
self.assertListEqual([], self.kbucket.peers)
|
|
|
|
# add the peer
|
|
self.kbucket.add_peer(peer)
|
|
self.assertListEqual([peer], self.kbucket.peers)
|
|
|
|
# re-add it
|
|
self.kbucket.add_peer(peer)
|
|
self.assertListEqual([peer], self.kbucket.peers)
|
|
self.assertEqual(self.kbucket.peers[0].udp_port, 4444)
|
|
|
|
# add a new peer object with the same id and address but a different port
|
|
self.kbucket.add_peer(peer_update2)
|
|
self.assertListEqual([peer_update2], self.kbucket.peers)
|
|
self.assertEqual(self.kbucket.peers[0].udp_port, 4445)
|
|
|
|
# modify the peer object to have a different port
|
|
peer_update2.udp_port = 4444
|
|
self.kbucket.add_peer(peer_update2)
|
|
self.assertListEqual([peer_update2], self.kbucket.peers)
|
|
self.assertEqual(self.kbucket.peers[0].udp_port, 4444)
|
|
|
|
self.kbucket.peers.clear()
|
|
|
|
# Test if contacts can be added to empty list
|
|
# Add k contacts to bucket
|
|
for i in range(constants.k):
|
|
peer = self.peer_manager.get_kademlia_peer(generate_id(), next(self.address_generator), 4444)
|
|
self.assertTrue(self.kbucket.add_peer(peer))
|
|
self.assertEqual(peer, self.kbucket.peers[i])
|
|
|
|
# Test if contact is not added to full list
|
|
peer = self.peer_manager.get_kademlia_peer(generate_id(), next(self.address_generator), 4444)
|
|
self.assertFalse(self.kbucket.add_peer(peer))
|
|
|
|
# Test if an existing contact is updated correctly if added again
|
|
existing_peer = self.kbucket.peers[0]
|
|
self.assertTrue(self.kbucket.add_peer(existing_peer))
|
|
self.assertEqual(existing_peer, self.kbucket.peers[-1])
|
|
|
|
# def testGetContacts(self):
|
|
# # try and get 2 contacts from empty list
|
|
# result = self.kbucket.getContacts(2)
|
|
# self.assertFalse(len(result) != 0, "Returned list should be empty; returned list length: %d" %
|
|
# (len(result)))
|
|
#
|
|
# # Add k-2 contacts
|
|
# node_ids = []
|
|
# if constants.k >= 2:
|
|
# for i in range(constants.k-2):
|
|
# node_ids.append(generate_id())
|
|
# tmpContact = self.contact_manager.make_contact(node_ids[-1], next(self.address_generator), 4444, 0,
|
|
# None)
|
|
# self.kbucket.addContact(tmpContact)
|
|
# else:
|
|
# # add k contacts
|
|
# for i in range(constants.k):
|
|
# node_ids.append(generate_id())
|
|
# tmpContact = self.contact_manager.make_contact(node_ids[-1], next(self.address_generator), 4444, 0,
|
|
# None)
|
|
# self.kbucket.addContact(tmpContact)
|
|
#
|
|
# # try to get too many contacts
|
|
# # requested count greater than bucket size; should return at most k contacts
|
|
# contacts = self.kbucket.getContacts(constants.k+3)
|
|
# self.assertTrue(len(contacts) <= constants.k,
|
|
# 'Returned list should not have more than k entries!')
|
|
#
|
|
# # verify returned contacts in list
|
|
# for node_id, i in zip(node_ids, range(constants.k-2)):
|
|
# self.assertFalse(self.kbucket._contacts[i].id != node_id,
|
|
# "Contact in position %s not same as added contact" % (str(i)))
|
|
#
|
|
# # try to get too many contacts
|
|
# # requested count one greater than number of contacts
|
|
# if constants.k >= 2:
|
|
# result = self.kbucket.getContacts(constants.k-1)
|
|
# self.assertFalse(len(result) != constants.k-2,
|
|
# "Too many contacts in returned list %s - should be %s" %
|
|
# (len(result), constants.k-2))
|
|
# else:
|
|
# result = self.kbucket.getContacts(constants.k-1)
|
|
# # if the count is <= 0, it should return all of it's contats
|
|
# self.assertFalse(len(result) != constants.k,
|
|
# "Too many contacts in returned list %s - should be %s" %
|
|
# (len(result), constants.k-2))
|
|
# result = self.kbucket.getContacts(constants.k-3)
|
|
# self.assertFalse(len(result) != constants.k-3,
|
|
# "Too many contacts in returned list %s - should be %s" %
|
|
# (len(result), constants.k-3))
|
|
|
|
def test_remove_peer(self):
|
|
# try remove contact from empty list
|
|
peer = self.peer_manager.get_kademlia_peer(generate_id(), next(self.address_generator), 4444)
|
|
self.assertRaises(ValueError, self.kbucket.remove_peer, peer)
|
|
|
|
added = []
|
|
# Add couple contacts
|
|
for i in range(constants.k-2):
|
|
peer = self.peer_manager.get_kademlia_peer(generate_id(), next(self.address_generator), 4444)
|
|
self.assertTrue(self.kbucket.add_peer(peer))
|
|
added.append(peer)
|
|
|
|
while added:
|
|
peer = added.pop()
|
|
self.assertIn(peer, self.kbucket.peers)
|
|
self.kbucket.remove_peer(peer)
|
|
self.assertNotIn(peer, self.kbucket.peers)
|