forked from LBRYCommunity/lbry-sdk
168 lines
7.1 KiB
Python
168 lines
7.1 KiB
Python
import time
|
|
import unittest
|
|
import twisted.internet.selectreactor
|
|
|
|
import lbrynet.dht.protocol
|
|
import lbrynet.dht.contact
|
|
import lbrynet.dht.constants
|
|
import lbrynet.dht.msgtypes
|
|
from lbrynet.dht.error import TimeoutError
|
|
from lbrynet.dht.node import Node, rpcmethod
|
|
|
|
|
|
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 = Node(node_id='1' * 48, udpPort=9182, externalIP="127.0.0.1")
|
|
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 """
|
|
|
|
@rpcmethod
|
|
def fake_ping(*args, **kwargs):
|
|
time.sleep(lbrynet.dht.constants.rpcTimeout + 1)
|
|
return 'pong'
|
|
|
|
real_ping = self.node.ping
|
|
real_timeout = lbrynet.dht.constants.rpcTimeout
|
|
real_attempts = lbrynet.dht.constants.rpcAttempts
|
|
lbrynet.dht.constants.rpcAttempts = 1
|
|
lbrynet.dht.constants.rpcTimeout = 1
|
|
self.node.ping = fake_ping
|
|
deadContact = lbrynet.dht.contact.Contact('2' * 48, '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)')
|
|
lbrynet.dht.protocol.reactor.listenUDP(9182, self.protocol)
|
|
|
|
# Run the PING RPC (which should raise a timeout error)
|
|
df = self.protocol.sendRPC(deadContact, 'ping', {})
|
|
|
|
def check_timeout(err):
|
|
self.assertEqual(type(err), TimeoutError)
|
|
|
|
df.addErrback(check_timeout)
|
|
|
|
def reset_values():
|
|
self.node.ping = real_ping
|
|
lbrynet.dht.constants.rpcTimeout = real_timeout
|
|
lbrynet.dht.constants.rpcAttempts = real_attempts
|
|
|
|
# See if the contact was removed due to the timeout
|
|
def check_removed_contact():
|
|
self.failIf(deadContact in self.node.contacts,
|
|
'Contact was not removed after RPC timeout; check exception types.')
|
|
|
|
df.addCallback(lambda _: reset_values())
|
|
|
|
# Stop the reactor if a result arrives (timeout or not)
|
|
df.addBoth(lambda _: lbrynet.dht.protocol.reactor.stop())
|
|
df.addCallback(lambda _: check_removed_contact())
|
|
lbrynet.dht.protocol.reactor.run()
|
|
|
|
def testRPCRequest(self):
|
|
""" Tests if a valid RPC request is executed and responded to correctly """
|
|
remoteContact = lbrynet.dht.contact.Contact('2' * 48, '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('2' * 48, '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.not_a_rpc_function()
|
|
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('2' * 48, '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!')
|