forked from LBRYCommunity/lbry-sdk
200 lines
8.2 KiB
Python
200 lines
8.2 KiB
Python
#!/usr/bin/env python
|
|
#
|
|
# This library is free software, distributed under the terms of
|
|
# the GNU Lesser General Public License Version 3, or any later version.
|
|
# See the COPYING file included in this archive
|
|
|
|
import time
|
|
import unittest
|
|
|
|
from twisted.internet import defer
|
|
from twisted.python import failure
|
|
import twisted.internet.selectreactor
|
|
from twisted.internet.protocol import DatagramProtocol
|
|
|
|
import lbrynet.dht.protocol
|
|
import lbrynet.dht.contact
|
|
import lbrynet.dht.constants
|
|
import lbrynet.dht.msgtypes
|
|
from lbrynet.dht.node import rpcmethod
|
|
|
|
|
|
class FakeNode(object):
|
|
""" A fake node object implementing some RPC and non-RPC methods to
|
|
test the Kademlia protocol's behaviour
|
|
"""
|
|
def __init__(self, id):
|
|
self.id = id
|
|
self.contacts = []
|
|
|
|
@rpcmethod
|
|
def ping(self):
|
|
return 'pong'
|
|
|
|
def pingNoRPC(self):
|
|
return 'pong'
|
|
|
|
@rpcmethod
|
|
def echo(self, value):
|
|
return value
|
|
|
|
def addContact(self, contact):
|
|
self.contacts.append(contact)
|
|
|
|
def removeContact(self, contact):
|
|
self.contacts.remove(contact)
|
|
|
|
def indirectPingContact(self, protocol, contact):
|
|
""" Pings the given contact (using the specified KademliaProtocol
|
|
object, not the direct Contact API), and removes the contact
|
|
on a timeout """
|
|
df = protocol.sendRPC(contact, 'ping', {})
|
|
def handleError(f):
|
|
if f.check(lbrynet.dht.protocol.TimeoutError):
|
|
self.removeContact(contact)
|
|
return f
|
|
else:
|
|
# This is some other error
|
|
return f
|
|
df.addErrback(handleError)
|
|
return df
|
|
|
|
class ClientDatagramProtocol(lbrynet.dht.protocol.KademliaProtocol):
|
|
data = ''
|
|
msgID = ''
|
|
destination = ('127.0.0.1', 9182)
|
|
|
|
def __init__(self):
|
|
lbrynet.dht.protocol.KademliaProtocol.__init__(self, None)
|
|
|
|
def startProtocol(self):
|
|
self.sendDatagram()
|
|
|
|
def sendDatagram(self):
|
|
if len(self.data):
|
|
self._send(self.data, self.msgID, self.destination)
|
|
|
|
|
|
|
|
|
|
class KademliaProtocolTest(unittest.TestCase):
|
|
""" Test case for the Protocol class """
|
|
def setUp(self):
|
|
del lbrynet.dht.protocol.reactor
|
|
lbrynet.dht.protocol.reactor = twisted.internet.selectreactor.SelectReactor()
|
|
self.node = FakeNode('node1')
|
|
self.protocol = lbrynet.dht.protocol.KademliaProtocol(self.node)
|
|
|
|
def testReactor(self):
|
|
""" Tests if the reactor can start/stop the protocol correctly """
|
|
lbrynet.dht.protocol.reactor.listenUDP(0, self.protocol)
|
|
lbrynet.dht.protocol.reactor.callLater(0, lbrynet.dht.protocol.reactor.stop)
|
|
lbrynet.dht.protocol.reactor.run()
|
|
|
|
def testRPCTimeout(self):
|
|
""" Tests if a RPC message sent to a dead remote node times out correctly """
|
|
deadContact = lbrynet.dht.contact.Contact('node2', '127.0.0.1', 9182, self.protocol)
|
|
self.node.addContact(deadContact)
|
|
# Make sure the contact was added
|
|
self.failIf(deadContact not in self.node.contacts, 'Contact not added to fake node (error in test code)')
|
|
# Set the timeout to 0 for testing
|
|
tempTimeout = lbrynet.dht.constants.rpcTimeout
|
|
lbrynet.dht.constants.rpcTimeout = 0
|
|
lbrynet.dht.protocol.reactor.listenUDP(0, self.protocol)
|
|
# Run the PING RPC (which should timeout)
|
|
df = self.node.indirectPingContact(self.protocol, deadContact)
|
|
# Stop the reactor if a result arrives (timeout or not)
|
|
df.addBoth(lambda _: lbrynet.dht.protocol.reactor.stop())
|
|
lbrynet.dht.protocol.reactor.run()
|
|
# See if the contact was removed due to the timeout
|
|
self.failIf(deadContact in self.node.contacts, 'Contact was not removed after RPC timeout; check exception types.')
|
|
# Restore the global timeout
|
|
lbrynet.dht.constants.rpcTimeout = tempTimeout
|
|
|
|
def testRPCRequest(self):
|
|
""" Tests if a valid RPC request is executed and responded to correctly """
|
|
remoteContact = lbrynet.dht.contact.Contact('node2', '127.0.0.1', 9182, self.protocol)
|
|
self.node.addContact(remoteContact)
|
|
self.error = None
|
|
def handleError(f):
|
|
self.error = 'An RPC error occurred: %s' % f.getErrorMessage()
|
|
def handleResult(result):
|
|
expectedResult = 'pong'
|
|
if result != expectedResult:
|
|
self.error = 'Result from RPC is incorrect; expected "%s", got "%s"' % (expectedResult, result)
|
|
# Publish the "local" node on the network
|
|
lbrynet.dht.protocol.reactor.listenUDP(9182, self.protocol)
|
|
# Simulate the RPC
|
|
df = remoteContact.ping()
|
|
df.addCallback(handleResult)
|
|
df.addErrback(handleError)
|
|
df.addBoth(lambda _: lbrynet.dht.protocol.reactor.stop())
|
|
lbrynet.dht.protocol.reactor.run()
|
|
self.failIf(self.error, self.error)
|
|
# The list of sent RPC messages should be empty at this stage
|
|
self.failUnlessEqual(len(self.protocol._sentMessages), 0, 'The protocol is still waiting for a RPC result, but the transaction is already done!')
|
|
|
|
def testRPCAccess(self):
|
|
""" Tests invalid RPC requests
|
|
|
|
Verifies that a RPC request for an existing but unpublished
|
|
method is denied, and that the associated (remote) exception gets
|
|
raised locally """
|
|
remoteContact = lbrynet.dht.contact.Contact('node2', '127.0.0.1', 9182, self.protocol)
|
|
self.node.addContact(remoteContact)
|
|
self.error = None
|
|
def handleError(f):
|
|
try:
|
|
f.raiseException()
|
|
except AttributeError, e:
|
|
# This is the expected outcome since the remote node did not publish the method
|
|
self.error = None
|
|
except Exception, e:
|
|
self.error = 'The remote method failed, but the wrong exception was raised; expected AttributeError, got %s' % type(e)
|
|
|
|
def handleResult(result):
|
|
self.error = 'The remote method executed successfully, returning: "%s"; this RPC should not have been allowed.' % result
|
|
# Publish the "local" node on the network
|
|
lbrynet.dht.protocol.reactor.listenUDP(9182, self.protocol)
|
|
# Simulate the RPC
|
|
df = remoteContact.pingNoRPC()
|
|
df.addCallback(handleResult)
|
|
df.addErrback(handleError)
|
|
df.addBoth(lambda _: lbrynet.dht.protocol.reactor.stop())
|
|
lbrynet.dht.protocol.reactor.run()
|
|
self.failIf(self.error, self.error)
|
|
# The list of sent RPC messages should be empty at this stage
|
|
self.failUnlessEqual(len(self.protocol._sentMessages), 0, 'The protocol is still waiting for a RPC result, but the transaction is already done!')
|
|
|
|
def testRPCRequestArgs(self):
|
|
""" Tests if an RPC requiring arguments is executed correctly """
|
|
remoteContact = lbrynet.dht.contact.Contact('node2', '127.0.0.1', 9182, self.protocol)
|
|
self.node.addContact(remoteContact)
|
|
self.error = None
|
|
def handleError(f):
|
|
self.error = 'An RPC error occurred: %s' % f.getErrorMessage()
|
|
def handleResult(result):
|
|
expectedResult = 'This should be returned.'
|
|
if result != 'This should be returned.':
|
|
self.error = 'Result from RPC is incorrect; expected "%s", got "%s"' % (expectedResult, result)
|
|
# Publish the "local" node on the network
|
|
lbrynet.dht.protocol.reactor.listenUDP(9182, self.protocol)
|
|
# Simulate the RPC
|
|
df = remoteContact.echo('This should be returned.')
|
|
df.addCallback(handleResult)
|
|
df.addErrback(handleError)
|
|
df.addBoth(lambda _: lbrynet.dht.protocol.reactor.stop())
|
|
lbrynet.dht.protocol.reactor.run()
|
|
self.failIf(self.error, self.error)
|
|
# The list of sent RPC messages should be empty at this stage
|
|
self.failUnlessEqual(len(self.protocol._sentMessages), 0, 'The protocol is still waiting for a RPC result, but the transaction is already done!')
|
|
|
|
|
|
def suite():
|
|
suite = unittest.TestSuite()
|
|
suite.addTest(unittest.makeSuite(KademliaProtocolTest))
|
|
return suite
|
|
|
|
if __name__ == '__main__':
|
|
# If this module is executed from the commandline, run all its tests
|
|
unittest.TextTestRunner().run(suite())
|