diff options
Diffstat (limited to 'test')
| -rw-r--r-- | test/mitmproxy/proxy/test_server.py | 34 | ||||
| -rw-r--r-- | test/mitmproxy/tservers.py | 8 |
2 files changed, 23 insertions, 19 deletions
diff --git a/test/mitmproxy/proxy/test_server.py b/test/mitmproxy/proxy/test_server.py index 4cfaa523..9e276294 100644 --- a/test/mitmproxy/proxy/test_server.py +++ b/test/mitmproxy/proxy/test_server.py @@ -276,10 +276,9 @@ class TestHTTP(tservers.HTTPProxyTest, CommonMixin): s = script.Script( tutils.test_data.path("mitmproxy/data/addonscripts/stream_modify.py") ) - self.master.addons.add(s) + self.set_addons(s) d = self.pathod('200:b"foo"') assert d.content == b"bar" - self.master.addons.remove(s) def test_first_line_rewrite(self): """ @@ -583,12 +582,11 @@ class TestTransparent(tservers.TransparentProxyTest, CommonMixin, TcpMixin): s = script.Script( tutils.test_data.path("mitmproxy/data/addonscripts/tcp_stream_modify.py") ) - self.master.addons.add(s) + self.set_addons(s) self._tcpproxy_on() d = self.pathod('200:b"foo"') self._tcpproxy_off() assert d.content == b"bar" - self.master.addons.remove(s) class TestTransparentSSL(tservers.TransparentProxyTest, CommonMixin, TcpMixin): @@ -731,7 +729,7 @@ class TestRedirectRequest(tservers.HTTPProxyTest): This test verifies that the original destination is restored for the third request. """ - self.proxy.tmaster.addons.add(ARedirectRequest(self.server2.port)) + self.set_addons(ARedirectRequest(self.server2.port)) p = self.pathoc() with p.connect(): @@ -770,7 +768,7 @@ class AStreamRequest: class TestStreamRequest(tservers.HTTPProxyTest): def test_stream_simple(self): - self.proxy.tmaster.addons.add(AStreamRequest()) + self.set_addons(AStreamRequest()) p = self.pathoc() with p.connect(): # a request with 100k of data but without content-length @@ -779,7 +777,7 @@ class TestStreamRequest(tservers.HTTPProxyTest): assert len(r1.content) > 100000 def test_stream_multiple(self): - self.proxy.tmaster.addons.add(AStreamRequest()) + self.set_addons(AStreamRequest()) p = self.pathoc() with p.connect(): # simple request with streaming turned on @@ -791,7 +789,7 @@ class TestStreamRequest(tservers.HTTPProxyTest): assert r1.status_code == 201 def test_stream_chunked(self): - self.proxy.tmaster.addons.add(AStreamRequest()) + self.set_addons(AStreamRequest()) connection = socket.socket(socket.AF_INET, socket.SOCK_STREAM) connection.connect(("127.0.0.1", self.proxy.port)) fconn = connection.makefile("rb") @@ -820,7 +818,7 @@ class AFakeResponse: class TestFakeResponse(tservers.HTTPProxyTest): def test_fake(self): - self.proxy.tmaster.addons.add(AFakeResponse()) + self.set_addons(AFakeResponse()) f = self.pathod("200") assert "header-response" in f.headers @@ -836,7 +834,7 @@ class TestServerConnect(tservers.HTTPProxyTest): def test_unnecessary_serverconnect(self): """A replayed/fake response with no upstream_cert should not connect to an upstream server""" - self.proxy.tmaster.addons.add(AFakeResponse()) + self.set_addons(AFakeResponse()) assert self.pathod("200").status_code == 200 assert not self.proxy.tmaster.has_log("serverconnect") @@ -849,7 +847,7 @@ class AKillRequest: class TestKillRequest(tservers.HTTPProxyTest): def test_kill(self): - self.proxy.tmaster.addons.add(AKillRequest()) + self.set_addons(AKillRequest()) with pytest.raises(exceptions.HttpReadDisconnect): self.pathod("200") # Nothing should have hit the server @@ -863,7 +861,7 @@ class AKillResponse: class TestKillResponse(tservers.HTTPProxyTest): def test_kill(self): - self.proxy.tmaster.addons.add(AKillResponse()) + self.set_addons(AKillResponse()) with pytest.raises(exceptions.HttpReadDisconnect): self.pathod("200") # The server should have seen a request @@ -886,7 +884,7 @@ class AIncomplete: class TestIncompleteResponse(tservers.HTTPProxyTest): def test_incomplete(self): - self.proxy.tmaster.addons.add(AIncomplete()) + self.set_addons(AIncomplete()) assert self.pathod("200").status_code == 502 @@ -969,7 +967,7 @@ class TestUpstreamProxySSL( def test_change_upstream_proxy_connect(self): # skip chain[0]. - self.proxy.tmaster.addons.add( + self.set_addons( UpstreamProxyChanger( ("127.0.0.1", self.chain[1].port) ) @@ -988,8 +986,8 @@ class TestUpstreamProxySSL( Client <- HTTPS -> Proxy <- HTTP -> Proxy <- HTTPS -> Server """ - self.proxy.tmaster.addons.add(RewriteToHttp()) - self.chain[1].tmaster.addons.add(RewriteToHttps()) + self.set_addons(RewriteToHttp()) + self.set_addons(RewriteToHttps()) p = self.pathoc() with p.connect(): resp = p.request("get:'/p/418'") @@ -1063,8 +1061,8 @@ class TestProxyChainingSSLReconnect(tservers.HTTPUpstreamProxyTest): http1obj.server_conn.wfile.write(headers) http1obj.server_conn.wfile.flush() - self.chain[0].tmaster.addons.add(RequestKiller([1, 2])) - self.chain[1].tmaster.addons.add(RequestKiller([1])) + self.chain[0].set_addons(RequestKiller([1, 2])) + self.chain[1].set_addons(RequestKiller([1])) p = self.pathoc() with p.connect(): diff --git a/test/mitmproxy/tservers.py b/test/mitmproxy/tservers.py index 9e6cef97..f3270e5c 100644 --- a/test/mitmproxy/tservers.py +++ b/test/mitmproxy/tservers.py @@ -92,6 +92,7 @@ class ProxyThread(threading.Thread): self.masterclass = masterclass self.options = options self.tmaster = None + self.event_loop = None controller.should_exit = False @property @@ -106,7 +107,8 @@ class ProxyThread(threading.Thread): self.tmaster.shutdown() def run(self): - asyncio.set_event_loop(asyncio.new_event_loop()) + self.event_loop = asyncio.new_event_loop() + asyncio.set_event_loop(self.event_loop) self.tmaster = self.masterclass(self.options) self.tmaster.addons.add(core.Core()) self.name = "ProxyThread (%s:%s)" % ( @@ -177,6 +179,10 @@ class ProxyTestBase: ssl_insecure=True, ) + def set_addons(self, *addons): + self.proxy.tmaster.reset(addons) + self.proxy.tmaster.addons.trigger("tick") + def addons(self): """ Can be over-ridden to add a standard set of addons to tests. |