lbry-sdk/lbrynet/tests/dht/test_protocol.py

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!')