From b0db225b9889b6bea55080663d5d4deecc2ad300 Mon Sep 17 00:00:00 2001
From: Justin Holmes <justin@justinholmes.com>
Date: Thu, 31 Aug 2017 14:36:22 -0700
Subject: [PATCH] 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           |  4 +++-
 kademlia/tests/test_server.py | 43 +++++++++++++++++++++++++++++++++++
 2 files changed, 46 insertions(+), 1 deletion(-)
 create mode 100644 kademlia/tests/test_server.py

diff --git a/kademlia/network.py b/kademlia/network.py
index 5530284..ad9a47e 100644
--- a/kademlia/network.py
+++ b/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)
diff --git a/kademlia/tests/test_server.py b/kademlia/tests/test_server.py
new file mode 100644
index 0000000..4688eac
--- /dev/null
+++ b/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()
\ No newline at end of file