lbry-sdk/tests/dht/testProtocol.py

#!/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())