Extract TCP test utilities into netlib.test

1 files changed, 67 insertions, 0 deletions

diff --git a/netlib/test.py b/netlib/test.py
new file mode 100644
index 00000000..2f72f979
--- /dev/null
+++ b/netlib/test.py
@@ -0,0 +1,67 @@
+import threading, Queue, cStringIO
+import tcp
+
+class ServerThread(threading.Thread):
+    def __init__(self, server):
+        self.server = server
+        threading.Thread.__init__(self)
+
+    def run(self):
+        self.server.serve_forever()
+
+    def shutdown(self):
+        self.server.shutdown()
+
+
+class ServerTestBase:
+    @classmethod
+    def setupAll(cls):
+        cls.q = Queue.Queue()
+        s = cls.makeserver()
+        cls.port = s.port
+        cls.server = ServerThread(s)
+        cls.server.start()
+
+    @classmethod
+    def teardownAll(cls):
+        cls.server.shutdown()
+
+
+    @property
+    def last_handler(self):
+        return self.server.server.last_handler
+
+
+class TServer(tcp.TCPServer):
+    def __init__(self, ssl, q, handler_klass, addr=("127.0.0.1", 0)):
+        """
+            ssl: A {cert, key, v3_only} dict.
+        """
+        tcp.TCPServer.__init__(self, addr)
+        self.ssl, self.q = ssl, q
+        self.handler_klass = handler_klass
+        self.last_handler = None
+
+    def handle_connection(self, request, client_address):
+        h = self.handler_klass(request, client_address, self)
+        self.last_handler = h
+        if self.ssl:
+            if self.ssl["v3_only"]:
+                method = tcp.SSLv3_METHOD
+                options = tcp.OP_NO_SSLv2|tcp.OP_NO_TLSv1
+            else:
+                method = tcp.SSLv23_METHOD
+                options = None
+            h.convert_to_ssl(
+                self.ssl["cert"],
+                self.ssl["key"],
+                method = method,
+                options = options,
+            )
+        h.handle()
+        h.finish()
+
+    def handle_error(self, request, client_address):
+        s = cStringIO.StringIO()
+        tcp.TCPServer.handle_error(self, request, client_address, s)
+        self.q.put(s.getvalue())