Significantly simplify server connection handling, and test.

3 files changed, 38 insertions, 48 deletions

diff --git a/test/test_proxy.py b/test/test_proxy.py
index b575a1d0..3995b393 100644
--- a/test/test_proxy.py
+++ b/test/test_proxy.py
@@ -39,8 +39,8 @@ class TestServerConnection:
         self.d.shutdown()
 
     def test_simple(self):
-        sc = proxy.ServerConnection(proxy.ProxyConfig(), self.d.IFACE, self.d.port)
-        sc.connect("http", "host.com")
+        sc = proxy.ServerConnection(proxy.ProxyConfig(), "http", self.d.IFACE, self.d.port, "host.com")
+        sc.connect()
         r = tutils.treq()
         r.path = "/p/200:da"
         sc.send(r)
@@ -53,36 +53,9 @@ class TestServerConnection:
         sc.terminate()
 
     def test_terminate_error(self):
-        sc = proxy.ServerConnection(proxy.ProxyConfig(), self.d.IFACE, self.d.port)
-        sc.connect("http", "host.com")
+        sc = proxy.ServerConnection(proxy.ProxyConfig(), "http", self.d.IFACE, self.d.port, "host.com")
+        sc.connect()
         sc.connection = mock.Mock()
         sc.connection.close = mock.Mock(side_effect=IOError)
         sc.terminate()
 
-
-
-def _dummysc(config, host, port):
-    return mock.MagicMock(config=config, host=host, port=port)
-
-
-def _errsc(config, host, port):
-    m = mock.MagicMock(config=config, host=host, port=port)
-    m.connect = mock.MagicMock(side_effect=tcp.NetLibError())
-    return m
-
-
-class TestServerConnectionPool:
-    @mock.patch("libmproxy.proxy.ServerConnection", _dummysc)
-    def test_pooling(self):
-        p = proxy.ServerConnectionPool(proxy.ProxyConfig())
-        c = p.get_connection("http", "localhost", 80, "localhost")
-        c2 = p.get_connection("http", "localhost", 80, "localhost")
-        assert c is c2
-        c3 = p.get_connection("http", "foo", 80, "localhost")
-        assert not c is c3
-
-    @mock.patch("libmproxy.proxy.ServerConnection", _errsc)
-    def test_connection_error(self):
-        p = proxy.ServerConnectionPool(proxy.ProxyConfig())
-        tutils.raises("502", p.get_connection, "http", "localhost", 80, "localhost")
-
diff --git a/test/test_server.py b/test/test_server.py
index 924b63b7..f93ddbb3 100644
--- a/test/test_server.py
+++ b/test/test_server.py
@@ -85,7 +85,7 @@ class TestHTTP(tservers.HTTPProxTest, SanityMixin):
     def test_connection_close(self):
         # Add a body, so we have a content-length header, which combined with
         # HTTP1.1 means the connection is kept alive.
-        response = '%s/p/200:b@1'%self.urlbase
+        response = '%s/p/200:b@1'%self.server.urlbase
 
         # Lets sanity check that the connection does indeed stay open by
         # issuing two requests over the same connection
@@ -99,7 +99,7 @@ class TestHTTP(tservers.HTTPProxTest, SanityMixin):
         tutils.raises("disconnect", p.request, "get:'%s'"%response)
 
     def test_reconnect(self):
-        req = "get:'%s/p/200:b@1:da'"%self.urlbase
+        req = "get:'%s/p/200:b@1:da'"%self.server.urlbase
         p = self.pathoc()
         assert p.request(req)
         # Server has disconnected. Mitmproxy should detect this, and reconnect.
@@ -107,7 +107,7 @@ class TestHTTP(tservers.HTTPProxTest, SanityMixin):
         assert p.request(req)
 
         # However, if the server disconnects on our first try, it's an error.
-        req = "get:'%s/p/200:b@1:d0'"%self.urlbase
+        req = "get:'%s/p/200:b@1:d0'"%self.server.urlbase
         p = self.pathoc()
         tutils.raises("server disconnect", p.request, req)
 
@@ -118,13 +118,29 @@ class TestHTTP(tservers.HTTPProxTest, SanityMixin):
             m.side_effect = IOError("error!")
             tutils.raises("empty reply", self.pathod, "304")
 
+    def test_get_connection_switching(self):
+        def switched(l):
+            for i in l:
+                if "switching" in i:
+                    return True
+        req = "get:'%s/p/200:b@1'"
+        p = self.pathoc()
+        assert p.request(req%self.server.urlbase)
+        assert p.request(req%self.server2.urlbase)
+        assert switched(self.proxy.log)
+
+    def test_get_connection_err(self):
+        p = self.pathoc()
+        ret = p.request("get:'http://localhost:0'")
+        assert ret[1] == 502
+
 
 class TestHTTPS(tservers.HTTPProxTest, SanityMixin):
     ssl = True
     clientcerts = True
     def test_clientcert(self):
         f = self.pathod("304")
-        assert self.last_log()["request"]["clientcert"]["keyinfo"]
+        assert self.server.last_log()["request"]["clientcert"]["keyinfo"]
 
 
 class TestReverse(tservers.ReverseProxTest, SanityMixin):
@@ -211,7 +227,7 @@ class TestKillRequest(tservers.HTTPProxTest):
         p = self.pathoc()
         tutils.raises("empty reply", self.pathod, "200")
         # Nothing should have hit the server
-        assert not self.last_log()
+        assert not self.server.last_log()
 
 
 class MasterKillResponse(tservers.TestMaster):
@@ -225,5 +241,5 @@ class TestKillResponse(tservers.HTTPProxTest):
         p = self.pathoc()
         tutils.raises("empty reply", self.pathod, "200")
         # The server should have seen a request
-        assert self.last_log()
+        assert self.server.last_log()
 
diff --git a/test/tservers.py b/test/tservers.py
index 262536a7..9597dab4 100644
--- a/test/tservers.py
+++ b/test/tservers.py
@@ -28,6 +28,7 @@ class TestMaster(flow.FlowMaster):
         state = flow.State()
         flow.FlowMaster.__init__(self, s, state)
         self.testq = testq
+        self.log = []
 
     def handle_request(self, m):
         flow.FlowMaster.handle_request(self, m)
@@ -37,6 +38,10 @@ class TestMaster(flow.FlowMaster):
         flow.FlowMaster.handle_response(self, m)
         m.reply()
 
+    def handle_log(self, l):
+        self.log.append(l.msg)
+        l.reply()
+
 
 class ProxyThread(threading.Thread):
     def __init__(self, tmaster):
@@ -48,6 +53,10 @@ class ProxyThread(threading.Thread):
     def port(self):
         return self.tmaster.server.port
 
+    @property
+    def log(self):
+        return self.tmaster.log
+
     def run(self):
         self.tmaster.run()
 
@@ -61,6 +70,7 @@ class ProxTestBase:
     def setupAll(cls):
         cls.tqueue = Queue.Queue()
         cls.server = libpathod.test.Daemon(ssl=cls.ssl)
+        cls.server2 = libpathod.test.Daemon(ssl=cls.ssl)
         pconf = cls.get_proxy_config()
         config = proxy.ProxyConfig(
             cacert = tutils.test_data.path("data/serverkey.pem"),
@@ -78,6 +88,7 @@ class ProxTestBase:
     def teardownAll(cls):
         cls.proxy.shutdown()
         cls.server.shutdown()
+        cls.server2.shutdown()
 
     def setUp(self):
         self.master.state.clear()
@@ -95,16 +106,6 @@ class ProxTestBase:
             (self.scheme, ("127.0.0.1", self.proxy.port))
         )
 
-    @property
-    def urlbase(self):
-        """
-            The URL base for the server instance.
-        """
-        return self.server.urlbase
-
-    def last_log(self):
-        return self.server.last_log()
-
 
 class HTTPProxTest(ProxTestBase):
     ssl = None
@@ -129,7 +130,7 @@ class HTTPProxTest(ProxTestBase):
             Constructs a pathod request, with the appropriate base and proxy.
         """
         return hurl.get(
-            self.urlbase + "/p/" + spec,
+            self.server.urlbase + "/p/" + spec,
             proxy=self.proxies,
             validate_cert=False,
             #debug=hurl.utils.stdout_debug