move tservers helper

5 files changed, 35 insertions, 33 deletions

diff --git a/test/mitmproxy/test_protocol_http2.py b/test/mitmproxy/test_protocol_http2.py
index fc27cc3f..1da140d8 100644
--- a/test/mitmproxy/test_protocol_http2.py
+++ b/test/mitmproxy/test_protocol_http2.py
@@ -18,7 +18,7 @@ logging.getLogger("PIL.Image").setLevel(logging.WARNING)
 logging.getLogger("PIL.PngImagePlugin").setLevel(logging.WARNING)
 
 import netlib
-from netlib import tservers as netlib_tservers
+from ..netlib import tservers as netlib_tservers
 from netlib.utils import http2_read_raw_frame
 
 import h2
diff --git a/test/netlib/http/http2/test_connections.py b/test/netlib/http/http2/test_connections.py
index 8be127e4..c067d487 100644
--- a/test/netlib/http/http2/test_connections.py
+++ b/test/netlib/http/http2/test_connections.py
@@ -4,11 +4,12 @@ import codecs
 
 from hyperframe.frame import *
 
-from netlib import tcp, http, utils, tservers
+from netlib import tcp, http, utils
 from netlib.tutils import raises
 from netlib.exceptions import TcpDisconnect
 from netlib.http.http2.connections import HTTP2Protocol, TCPHandler
 
+from ... import tservers
 
 class TestTCPHandlerWrapper:
     def test_wrapped(self):
diff --git a/test/netlib/test_tcp.py b/test/netlib/test_tcp.py
index 8ae3aa51..e65a2e2f 100644
--- a/test/netlib/test_tcp.py
+++ b/test/netlib/test_tcp.py
@@ -10,10 +10,11 @@ import mock
 from OpenSSL import SSL
 import OpenSSL
 
-from netlib import tcp, certutils, tutils, tservers
+from netlib import tcp, certutils, tutils
 from netlib.exceptions import InvalidCertificateException, TcpReadIncomplete, TlsException, \
     TcpTimeout, TcpDisconnect, TcpException, NetlibException
 
+from . import tservers
 
 class EchoHandler(tcp.BaseHandler):
     sni = None
diff --git a/netlib/netlib/tservers.py b/test/netlib/tservers.py
index 44ef8063..569745e6 100644
--- a/netlib/netlib/tservers.py
+++ b/test/netlib/tservers.py
@@ -9,7 +9,7 @@ from netlib import tcp
 from netlib import tutils
 
 
-class ServerThread(threading.Thread):
+class _ServerThread(threading.Thread):
 
     def __init__(self, server):
         self.server = server
@@ -22,33 +22,7 @@ class ServerThread(threading.Thread):
         self.server.shutdown()
 
 
-class ServerTestBase(object):
-    ssl = None
-    handler = None
-    addr = ("localhost", 0)
-
-    @classmethod
-    def setup_class(cls):
-        cls.q = queue.Queue()
-        s = cls.makeserver()
-        cls.port = s.address.port
-        cls.server = ServerThread(s)
-        cls.server.start()
-
-    @classmethod
-    def makeserver(cls):
-        return TServer(cls.ssl, cls.q, cls.handler, cls.addr)
-
-    @classmethod
-    def teardown_class(cls):
-        cls.server.shutdown()
-
-    @property
-    def last_handler(self):
-        return self.server.server.last_handler
-
-
-class TServer(tcp.TCPServer):
+class _TServer(tcp.TCPServer):
 
     def __init__(self, ssl, q, handler_klass, addr):
         """
@@ -107,3 +81,29 @@ class TServer(tcp.TCPServer):
         s = StringIO()
         tcp.TCPServer.handle_error(self, connection, client_address, s)
         self.q.put(s.getvalue())
+
+
+class ServerTestBase(object):
+    ssl = None
+    handler = None
+    addr = ("localhost", 0)
+
+    @classmethod
+    def setup_class(cls):
+        cls.q = queue.Queue()
+        s = cls.makeserver()
+        cls.port = s.address.port
+        cls.server = _ServerThread(s)
+        cls.server.start()
+
+    @classmethod
+    def makeserver(cls):
+        return _TServer(cls.ssl, cls.q, cls.handler, cls.addr)
+
+    @classmethod
+    def teardown_class(cls):
+        cls.server.shutdown()
+
+    @property
+    def last_handler(self):
+        return self.server.server.last_handler
diff --git a/test/netlib/websockets/test_websockets.py b/test/netlib/websockets/test_websockets.py
index d53f0d83..a7d782a4 100644
--- a/test/netlib/websockets/test_websockets.py
+++ b/test/netlib/websockets/test_websockets.py
@@ -2,12 +2,12 @@ import os
 
 from netlib.http.http1 import read_response, read_request
 
-from netlib import tcp, websockets, http, tutils, tservers
+from netlib import tcp, websockets, http, tutils
 from netlib.http import status_codes
 from netlib.tutils import treq
-
 from netlib.exceptions import *
 
+from .. import tservers
 
 class WebSocketsEchoHandler(tcp.BaseHandler):