lbry-sdk/tests/unit/dht/test_node.py

import hashlib
import struct

from twisted.trial import unittest
from twisted.internet import defer
from lbrynet.dht.node import Node
from lbrynet.dht import constants
from lbrynet.core.utils import generate_id


class NodeIDTest(unittest.TestCase):

    def setUp(self):
        self.node = Node()

    def test_new_node_has_auto_created_id(self):
        self.assertEqual(type(self.node.node_id), bytes)
        self.assertEqual(len(self.node.node_id), 48)

    def test_uniqueness_and_length_of_generated_ids(self):
        previous_ids = []
        for i in range(100):
            new_id = self.node._generateID()
            self.assertNotIn(new_id, previous_ids, 'id at index {} not unique'.format(i))
            self.assertEqual(len(new_id), 48, 'id at index {} wrong length: {}'.format(i, len(new_id)))
            previous_ids.append(new_id)


class NodeDataTest(unittest.TestCase):
    """ Test case for the Node class's data-related functions """

    def setUp(self):
        h = hashlib.sha384()
        h.update(b'test')
        self.node = Node()
        self.contact = self.node.contact_manager.make_contact(
            h.digest(), '127.0.0.1', 12345, self.node._protocol)
        self.token = self.node.make_token(self.contact.compact_ip())
        self.cases = []
        for i in range(5):
            h.update(str(i).encode())
            self.cases.append((h.digest(), 5000+2*i))
            self.cases.append((h.digest(), 5001+2*i))

    @defer.inlineCallbacks
    def test_store(self):
        """ Tests if the node can store (and privately retrieve) some data """
        for key, port in self.cases:
            yield self.node.store(
                self.contact, key, self.token, port, self.contact.id, 0
            )
        for key, value in self.cases:
            expected_result = self.contact.compact_ip() + struct.pack('>H', value) + self.contact.id
            self.assertTrue(self.node._dataStore.hasPeersForBlob(key),
                            "Stored key not found in node's DataStore: '%s'" % key)
            self.assertTrue(expected_result in self.node._dataStore.getPeersForBlob(key),
                            "Stored val not found in node's DataStore: key:'%s' port:'%s' %s"
                            % (key, value, self.node._dataStore.getPeersForBlob(key)))


class NodeContactTest(unittest.TestCase):
    """ Test case for the Node class's contact management-related functions """
    def setUp(self):
        self.node = Node()

    @defer.inlineCallbacks
    def test_add_contact(self):
        """ Tests if a contact can be added and retrieved correctly """
        # Create the contact
        contact_id = generate_id(b'node1')
        contact = self.node.contact_manager.make_contact(contact_id, '127.0.0.1', 9182, self.node._protocol)
        # Now add it...
        yield self.node.addContact(contact)
        # ...and request the closest nodes to it using FIND_NODE
        closest_nodes = self.node._routingTable.findCloseNodes(contact_id, constants.k)
        self.assertEqual(len(closest_nodes), 1)
        self.assertIn(contact, closest_nodes)

    @defer.inlineCallbacks
    def test_add_self_as_contact(self):
        """ Tests the node's behaviour when attempting to add itself as a contact """
        # Create a contact with the same ID as the local node's ID
        contact = self.node.contact_manager.make_contact(self.node.node_id, '127.0.0.1', 9182, None)
        # Now try to add it
        yield self.node.addContact(contact)
        # ...and request the closest nodes to it using FIND_NODE
        closest_nodes = self.node._routingTable.findCloseNodes(self.node.node_id, constants.k)
        self.assertNotIn(contact, closest_nodes, 'Node added itself as a contact.')