import hashlib from twisted.trial import unittest from twisted.internet import defer from lbrynet.dht import constants from lbrynet.dht.routingtable import TreeRoutingTable from lbrynet.dht.contact import ContactManager from lbrynet.dht.distance import Distance class FakeRPCProtocol(object): """ Fake RPC protocol; allows lbrynet.dht.contact.Contact objects to "send" RPCs """ def sendRPC(self, *args, **kwargs): return defer.succeed(None) class TreeRoutingTableTest(unittest.TestCase): """ Test case for the RoutingTable class """ def setUp(self): h = hashlib.sha384() h.update('node1') self.contact_manager = ContactManager() self.nodeID = h.digest() self.protocol = FakeRPCProtocol() self.routingTable = TreeRoutingTable(self.nodeID) def testDistance(self): """ Test to see if distance method returns correct result""" # testList holds a couple 3-tuple (variable1, variable2, result) basicTestList = [(chr(170) * 48, chr(85) * 48, long((chr(255) * 48).encode('hex'), 16))] for test in basicTestList: result = Distance(test[0])(test[1]) self.failIf(result != test[2], 'Result of _distance() should be %s but %s returned' % (test[2], result)) @defer.inlineCallbacks def testAddContact(self): """ Tests if a contact can be added and retrieved correctly """ # Create the contact h = hashlib.sha384() h.update('node2') contactID = h.digest() contact = self.contact_manager.make_contact(contactID, '127.0.0.1', 9182, self.protocol) # Now add it... yield self.routingTable.addContact(contact) # ...and request the closest nodes to it (will retrieve it) closestNodes = self.routingTable.findCloseNodes(contactID, constants.k) self.failUnlessEqual(len(closestNodes), 1, 'Wrong amount of contacts returned; expected 1,' ' got %d' % len(closestNodes)) self.failUnless(contact in closestNodes, 'Added contact not found by issueing ' '_findCloseNodes()') @defer.inlineCallbacks def testGetContact(self): """ Tests if a specific existing contact can be retrieved correctly """ h = hashlib.sha384() h.update('node2') contactID = h.digest() contact = self.contact_manager.make_contact(contactID, '127.0.0.1', 9182, self.protocol) # Now add it... yield self.routingTable.addContact(contact) # ...and get it again sameContact = self.routingTable.getContact(contactID) self.failUnlessEqual(contact, sameContact, 'getContact() should return the same contact') @defer.inlineCallbacks def testAddParentNodeAsContact(self): """ Tests the routing table's behaviour when attempting to add its parent node as a contact """ # Create a contact with the same ID as the local node's ID contact = self.contact_manager.make_contact(self.nodeID, '127.0.0.1', 9182, self.protocol) # Now try to add it yield self.routingTable.addContact(contact) # ...and request the closest nodes to it using FIND_NODE closestNodes = self.routingTable.findCloseNodes(self.nodeID, constants.k) self.failIf(contact in closestNodes, 'Node added itself as a contact') @defer.inlineCallbacks def testRemoveContact(self): """ Tests contact removal """ # Create the contact h = hashlib.sha384() h.update('node2') contactID = h.digest() contact = self.contact_manager.make_contact(contactID, '127.0.0.1', 9182, self.protocol) # Now add it... yield self.routingTable.addContact(contact) # Verify addition self.failUnlessEqual(len(self.routingTable._buckets[0]), 1, 'Contact not added properly') # Now remove it self.routingTable.removeContact(contact) self.failUnlessEqual(len(self.routingTable._buckets[0]), 0, 'Contact not removed properly') @defer.inlineCallbacks def testSplitBucket(self): """ Tests if the the routing table correctly dynamically splits k-buckets """ self.failUnlessEqual(self.routingTable._buckets[0].rangeMax, 2**384, 'Initial k-bucket range should be 0 <= range < 2**384') # Add k contacts for i in range(constants.k): h = hashlib.sha384() h.update('remote node %d' % i) nodeID = h.digest() contact = self.contact_manager.make_contact(nodeID, '127.0.0.1', 9182, self.protocol) yield self.routingTable.addContact(contact) self.failUnlessEqual(len(self.routingTable._buckets), 1, 'Only k nodes have been added; the first k-bucket should now ' 'be full, but should not yet be split') # Now add 1 more contact h = hashlib.sha384() h.update('yet another remote node') nodeID = h.digest() contact = self.contact_manager.make_contact(nodeID, '127.0.0.1', 9182, self.protocol) yield self.routingTable.addContact(contact) self.failUnlessEqual(len(self.routingTable._buckets), 2, 'k+1 nodes have been added; the first k-bucket should have been ' 'split into two new buckets') self.failIfEqual(self.routingTable._buckets[0].rangeMax, 2**384, 'K-bucket was split, but its range was not properly adjusted') self.failUnlessEqual(self.routingTable._buckets[1].rangeMax, 2**384, 'K-bucket was split, but the second (new) bucket\'s ' 'max range was not set properly') self.failUnlessEqual(self.routingTable._buckets[0].rangeMax, self.routingTable._buckets[1].rangeMin, 'K-bucket was split, but the min/max ranges were ' 'not divided properly') @defer.inlineCallbacks def testFullSplit(self): """ Test that a bucket is not split if it is full, but the new contact is not closer than the kth closest contact """ self.routingTable._parentNodeID = 48 * chr(255) node_ids = [ "100000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000", "200000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000", "300000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000", "400000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000", "500000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000", "600000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000", "700000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000", "800000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000", "ff0000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000", "010000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000" ] # Add k contacts for nodeID in node_ids: # self.assertEquals(nodeID, node_ids[i].decode('hex')) contact = self.contact_manager.make_contact(nodeID.decode('hex'), '127.0.0.1', 9182, self.protocol) yield self.routingTable.addContact(contact) self.failUnlessEqual(len(self.routingTable._buckets), 2) self.failUnlessEqual(len(self.routingTable._buckets[0]._contacts), 8) self.failUnlessEqual(len(self.routingTable._buckets[1]._contacts), 2) # try adding a contact who is further from us than the k'th known contact nodeID = '020000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000' nodeID = nodeID.decode('hex') contact = self.contact_manager.make_contact(nodeID, '127.0.0.1', 9182, self.protocol) self.assertFalse(self.routingTable._shouldSplit(self.routingTable._kbucketIndex(contact.id), contact.id)) yield self.routingTable.addContact(contact) self.failUnlessEqual(len(self.routingTable._buckets), 2) self.failUnlessEqual(len(self.routingTable._buckets[0]._contacts), 8) self.failUnlessEqual(len(self.routingTable._buckets[1]._contacts), 2) self.failIf(contact in self.routingTable._buckets[0]._contacts) self.failIf(contact in self.routingTable._buckets[1]._contacts) # class KeyErrorFixedTest(unittest.TestCase): # """ Basic tests case for boolean operators on the Contact class """ # # def setUp(self): # own_id = (2 ** constants.key_bits) - 1 # # carefully chosen own_id. here's the logic # # we want a bunch of buckets (k+1, to be exact), and we want to make sure own_id # # is not in bucket 0. so we put own_id at the end so we can keep splitting by adding to the # # end # # self.table = lbrynet.dht.routingtable.OptimizedTreeRoutingTable(own_id) # # def fill_bucket(self, bucket_min): # bucket_size = lbrynet.dht.constants.k # for i in range(bucket_min, bucket_min + bucket_size): # self.table.addContact(lbrynet.dht.contact.Contact(long(i), '127.0.0.1', 9999, None)) # # def overflow_bucket(self, bucket_min): # bucket_size = lbrynet.dht.constants.k # self.fill_bucket(bucket_min) # self.table.addContact( # lbrynet.dht.contact.Contact(long(bucket_min + bucket_size + 1), # '127.0.0.1', 9999, None)) # # def testKeyError(self): # # # find middle, so we know where bucket will split # bucket_middle = self.table._buckets[0].rangeMax / 2 # # # fill last bucket # self.fill_bucket(self.table._buckets[0].rangeMax - lbrynet.dht.constants.k - 1) # # -1 in previous line because own_id is in last bucket # # # fill/overflow 7 more buckets # bucket_start = 0 # for i in range(0, lbrynet.dht.constants.k): # self.overflow_bucket(bucket_start) # bucket_start += bucket_middle / (2 ** i) # # # replacement cache now has k-1 entries. # # adding one more contact to bucket 0 used to cause a KeyError, but it should work # self.table.addContact( # lbrynet.dht.contact.Contact(long(lbrynet.dht.constants.k + 2), '127.0.0.1', 9999, None)) # # # import math # # print "" # # for i, bucket in enumerate(self.table._buckets): # # print "Bucket " + str(i) + " (2 ** " + str( # # math.log(bucket.rangeMin, 2) if bucket.rangeMin > 0 else 0) + " <= x < 2 ** "+str( # # math.log(bucket.rangeMax, 2)) + ")" # # for c in bucket.getContacts(): # # print " contact " + str(c.id) # # for key, bucket in self.table._replacementCache.iteritems(): # # print "Replacement Cache for Bucket " + str(key) # # for c in bucket: # # print " contact " + str(c.id)