forked from LBRYCommunity/lbry-sdk
89 lines
3.6 KiB
Python
89 lines
3.6 KiB
Python
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.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, f'id at index {i} not unique')
|
|
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.assertIn(expected_result, 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.')
|