fix bugs, fix tests

3 files changed, 26 insertions, 63 deletions

diff --git a/test/test_proxy.py b/test/test_proxy.py
index 9c01ab63..fac4a4f4 100644
--- a/test/test_proxy.py
+++ b/test/test_proxy.py
@@ -1,9 +1,8 @@
-import argparse
 from libmproxy import cmdline
 from libmproxy.proxy import ProxyConfig, process_proxy_options
 from libmproxy.proxy.connection import ServerConnection
 from libmproxy.proxy.primitives import ProxyError
-from libmproxy.proxy.server import DummyServer, ProxyServer, ConnectionHandler
+from libmproxy.proxy.server import DummyServer, ProxyServer, ConnectionHandler2
 import tutils
 from libpathod import test
 from netlib import http, tcp
@@ -175,8 +174,10 @@ class TestDummyServer:
 class TestConnectionHandler:
     def test_fatal_error(self):
         config = mock.Mock()
-        config.mode.get_upstream_server.side_effect = RuntimeError
-        c = ConnectionHandler(
+        root_layer = mock.Mock()
+        root_layer.side_effect = RuntimeError
+        config.mode.return_value = root_layer
+        c = ConnectionHandler2(
             config,
             mock.MagicMock(),
             ("127.0.0.1",
diff --git a/test/test_server.py b/test/test_server.py
index 1216a349..7b66c582 100644
--- a/test/test_server.py
+++ b/test/test_server.py
@@ -68,7 +68,7 @@ class CommonMixin:
         # SSL with the upstream proxy.
         rt = self.master.replay_request(l, block=True)
         assert not rt
-        if isinstance(self, tservers.HTTPUpstreamProxTest) and not self.ssl:
+        if isinstance(self, tservers.HTTPUpstreamProxTest):
             assert l.response.code == 502
         else:
             assert l.error
@@ -506,7 +506,7 @@ class TestTransparentSSL(tservers.TransparentProxTest, CommonMixin, TcpMixin):
         p = pathoc.Pathoc(("localhost", self.proxy.port), fp=None)
         p.connect()
         r = p.request("get:/")
-        assert r.status_code == 400
+        assert r.status_code == 502
 
 
 class TestProxy(tservers.HTTPProxTest):
@@ -724,9 +724,9 @@ class TestStreamRequest(tservers.HTTPProxTest):
         assert resp.headers["Transfer-Encoding"][0] == 'chunked'
         assert resp.status_code == 200
 
-        chunks = list(
-            content for _, content, _ in protocol.read_http_body_chunked(
-                resp.headers, None, "GET", 200, False))
+        chunks = list(protocol.read_http_body_chunked(
+            resp.headers, None, "GET", 200, False
+        ))
         assert chunks == ["this", "isatest", ""]
 
         connection.close()
@@ -959,6 +959,9 @@ class TestProxyChainingSSLReconnect(tservers.HTTPUpstreamProxTest):
 
         p = self.pathoc()
         req = p.request("get:'/p/418:b\"content\"'")
+        assert req.content == "content"
+        assert req.status_code == 418
+
         assert self.proxy.tmaster.state.flow_count() == 2  # CONNECT and request
         # CONNECT, failing request,
         assert self.chain[0].tmaster.state.flow_count() == 4
@@ -967,8 +970,7 @@ class TestProxyChainingSSLReconnect(tservers.HTTPUpstreamProxTest):
         assert self.chain[1].tmaster.state.flow_count() == 2
         # (doesn't store (repeated) CONNECTs from chain[0]
         #  as it is a regular proxy)
-        assert req.content == "content"
-        assert req.status_code == 418
+
 
         assert not self.chain[1].tmaster.state.flows[0].response  # killed
         assert self.chain[1].tmaster.state.flows[1].response
diff --git a/test/tservers.py b/test/tservers.py
index 43ebf2bb..dfd3f627 100644
--- a/test/tservers.py
+++ b/test/tservers.py
@@ -181,22 +181,24 @@ class TResolver:
     def original_addr(self, sock):
         return ("127.0.0.1", self.port)
 
-
 class TransparentProxTest(ProxTestBase):
     ssl = None
     resolver = TResolver
 
     @classmethod
-    @mock.patch("libmproxy.platform.resolver")
-    def setupAll(cls, _):
+    def setupAll(cls):
         super(TransparentProxTest, cls).setupAll()
-        if cls.ssl:
-            ports = [cls.server.port, cls.server2.port]
-        else:
-            ports = []
-        cls.config.mode = TransparentProxyMode(
-            cls.resolver(cls.server.port),
-            ports)
+
+        cls._resolver = mock.patch(
+            "libmproxy.platform.resolver",
+            new=lambda: cls.resolver(cls.server.port)
+        )
+        cls._resolver.start()
+
+    @classmethod
+    def teardownAll(cls):
+        cls._resolver.stop()
+        super(TransparentProxTest, cls).teardownAll()
 
     @classmethod
     def get_proxy_config(cls):
@@ -270,48 +272,6 @@ class SocksModeTest(HTTPProxTest):
         d["mode"] = "socks5"
         return d
 
-class SpoofModeTest(ProxTestBase):
-    ssl = None
-
-    @classmethod
-    def get_proxy_config(cls):
-        d = ProxTestBase.get_proxy_config()
-        d["upstream_server"] = None
-        d["mode"] = "spoof"
-        return d
-
-    def pathoc(self, sni=None):
-        """
-            Returns a connected Pathoc instance.
-        """
-        p = libpathod.pathoc.Pathoc(
-            ("localhost", self.proxy.port), ssl=self.ssl, sni=sni, fp=None
-        )
-        p.connect()
-        return p
-
-
-class SSLSpoofModeTest(ProxTestBase):
-    ssl = True
-
-    @classmethod
-    def get_proxy_config(cls):
-        d = ProxTestBase.get_proxy_config()
-        d["upstream_server"] = None
-        d["mode"] = "sslspoof"
-        d["spoofed_ssl_port"] = 443
-        return d
-
-    def pathoc(self, sni=None):
-        """
-            Returns a connected Pathoc instance.
-        """
-        p = libpathod.pathoc.Pathoc(
-            ("localhost", self.proxy.port), ssl=self.ssl, sni=sni, fp=None
-        )
-        p.connect()
-        return p
-
 
 class ChainProxTest(ProxTestBase):
     """