lbry-sdk/lbry/tests/unit/dht/protocol/test_kbucket.py

import struct
import asyncio
from lbry.utils import generate_id
from lbry.dht.protocol.routing_table import KBucket
from lbry.dht.peer import PeerManager, make_kademlia_peer
from lbry.dht import constants
from lbry.testcase import AsyncioTestCase


def address_generator(address=(1, 2, 3, 4)):
    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 = make_kademlia_peer(constants.generate_id(2), "1.2.3.4", udp_port=4444)
        peer_update2 = make_kademlia_peer(constants.generate_id(2), "1.2.3.4", 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 = make_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 = make_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 = make_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 = make_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)