Allowing swappable protocol_class for Server. (#32)

* Allowing swappable protocol_class for Server.

* Changed protocol_class to be an attribute on the Server class; added tests.

kademlia/network.py

@@ -20,6 +20,8 @@ class Server(object):
     to start listening as an active node on the network.
     """
 
+    protocol_class = KademliaProtocol
+
     def __init__(self, ksize=20, alpha=3, id=None, storage=None):
         """
         Create a server instance.  This will start listening on the given port.
@@ -52,7 +54,7 @@ class Server(object):
 
         Provide interface="::" to accept ipv6 address
         """
-        proto_factory = lambda: KademliaProtocol(self.node, self.storage, self.ksize)
+        proto_factory = lambda: self.protocol_class(self.node, self.storage, self.ksize)
         loop = asyncio.get_event_loop()
         listen = loop.create_datagram_endpoint(proto_factory, local_addr=(interface, port))
         self.transport, self.protocol = loop.run_until_complete(listen)

kademlia/tests/test_server.py

@@ -0,0 +1,43 @@
+import unittest
+
+from kademlia.network import Server
+from kademlia.protocol import KademliaProtocol
+
+
+class SwappableProtocolTests(unittest.TestCase):
+
+    def test_default_protocol(self):
+        """
+        An ordinary Server object will initially not have a protocol, but will have a KademliaProtocol
+        object as its protocol after its listen() method is called.
+        """
+        server = Server()
+        self.assertIsNone(server.protocol)
+        server.listen(8469)
+        self.assertIsInstance(server.protocol, KademliaProtocol)
+        server.stop()
+
+    def test_custom_protocol(self):
+        """
+        A subclass of Server which overrides the protocol_class attribute will have an instance
+        of that class as its protocol after its listen() method is called.
+        """
+
+        # Make a custom Protocol and Server to go with hit.
+        class CoconutProtocol(KademliaProtocol):
+            pass
+
+        class HuskServer(Server):
+            protocol_class = CoconutProtocol
+
+        # An ordinary server does NOT have a CoconutProtocol as its protocol...
+        server = Server()
+        server.listen(8469)
+        self.assertNotIsInstance(server.protocol, CoconutProtocol)
+        server.stop()
+
+        # ...but our custom server does.
+        husk_server = HuskServer()
+        husk_server.listen(8469)
+        self.assertIsInstance(husk_server.protocol, CoconutProtocol)
+        server.stop()