diff options
Diffstat (limited to 'test')
| -rw-r--r-- | test/mitmproxy/test_controller.py | 31 | ||||
| -rw-r--r-- | test/mitmproxy/test_flow.py | 14 | ||||
| -rw-r--r-- | test/mitmproxy/test_server.py | 36 | ||||
| -rw-r--r-- | test/mitmproxy/tservers.py | 4 |
4 files changed, 48 insertions, 37 deletions
diff --git a/test/mitmproxy/test_controller.py b/test/mitmproxy/test_controller.py index f7bf615a..c9a8e2f4 100644 --- a/test/mitmproxy/test_controller.py +++ b/test/mitmproxy/test_controller.py @@ -2,7 +2,7 @@ from threading import Thread, Event from mock import Mock -from mitmproxy.controller import Reply, DummyReply, Channel, ServerThread, ServerMaster, Master +from mitmproxy import controller from six.moves import queue from mitmproxy.exceptions import Kill @@ -10,10 +10,15 @@ from mitmproxy.proxy import DummyServer from netlib.tutils import raises +class TMsg: + pass + + class TestMaster(object): def test_simple(self): - class DummyMaster(Master): + class DummyMaster(controller.Master): + @controller.handler def handle_panic(self, _): m.should_exit.set() @@ -23,14 +28,16 @@ class TestMaster(object): m = DummyMaster() assert not m.should_exit.is_set() - m.event_queue.put(("panic", 42)) + msg = TMsg() + msg.reply = controller.DummyReply() + m.event_queue.put(("panic", msg)) m.run() assert m.should_exit.is_set() class TestServerMaster(object): def test_simple(self): - m = ServerMaster() + m = controller.ServerMaster() s = DummyServer(None) m.add_server(s) m.start() @@ -42,7 +49,7 @@ class TestServerMaster(object): class TestServerThread(object): def test_simple(self): m = Mock() - t = ServerThread(m) + t = controller.ServerThread(m) t.run() assert m.serve_forever.called @@ -50,7 +57,7 @@ class TestServerThread(object): class TestChannel(object): def test_tell(self): q = queue.Queue() - channel = Channel(q, Event()) + channel = controller.Channel(q, Event()) m = Mock() channel.tell("test", m) assert q.get() == ("test", m) @@ -66,21 +73,21 @@ class TestChannel(object): Thread(target=reply).start() - channel = Channel(q, Event()) + channel = controller.Channel(q, Event()) assert channel.ask("test", Mock()) == 42 def test_ask_shutdown(self): q = queue.Queue() done = Event() done.set() - channel = Channel(q, done) + channel = controller.Channel(q, done) with raises(Kill): channel.ask("test", Mock()) class TestDummyReply(object): def test_simple(self): - reply = DummyReply() + reply = controller.DummyReply() assert not reply.acked reply() assert reply.acked @@ -88,18 +95,18 @@ class TestDummyReply(object): class TestReply(object): def test_simple(self): - reply = Reply(42) + reply = controller.Reply(42) assert not reply.acked reply("foo") assert reply.acked assert reply.q.get() == "foo" def test_default(self): - reply = Reply(42) + reply = controller.Reply(42) reply() assert reply.q.get() == 42 def test_reply_none(self): - reply = Reply(42) + reply = controller.Reply(42) reply(None) assert reply.q.get() is None diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py index 3e78a5c4..da8b8ddd 100644 --- a/test/mitmproxy/test_flow.py +++ b/test/mitmproxy/test_flow.py @@ -461,25 +461,20 @@ class TestFlow(object): fm = flow.FlowMaster(None, s) f = tutils.tflow() f.intercept(mock.Mock()) - assert not f.reply.acked f.kill(fm) - assert f.reply.acked + for i in s.view: + assert "killed" in str(i.error) def test_killall(self): s = flow.State() fm = flow.FlowMaster(None, s) f = tutils.tflow() - fm.handle_request(f) - - f = tutils.tflow() - fm.handle_request(f) + f.intercept(fm) - for i in s.view: - assert not i.reply.acked s.killall(fm) for i in s.view: - assert i.reply.acked + assert "killed" in str(i.error) def test_accept_intercept(self): f = tutils.tflow() @@ -866,7 +861,6 @@ class TestFlowMaster: f.response = HTTPResponse.wrap(netlib.tutils.tresp()) fm.handle_response(f) - assert not fm.handle_response(None) assert s.flow_count() == 1 fm.handle_clientdisconnect(f.client_conn) diff --git a/test/mitmproxy/test_server.py b/test/mitmproxy/test_server.py index 0701d52b..f4e7452f 100644 --- a/test/mitmproxy/test_server.py +++ b/test/mitmproxy/test_server.py @@ -12,6 +12,7 @@ from netlib.http import authentication, http1 from netlib.tutils import raises from pathod import pathoc, pathod +from mitmproxy import controller from mitmproxy.proxy.config import HostMatcher from mitmproxy.exceptions import Kill from mitmproxy.models import Error, HTTPResponse, HTTPFlow @@ -623,6 +624,7 @@ class TestProxySSL(tservers.HTTPProxyTest): class MasterRedirectRequest(tservers.TestMaster): redirect_port = None # Set by TestRedirectRequest + @controller.handler def handle_request(self, f): if f.request.path == "/p/201": @@ -636,6 +638,7 @@ class MasterRedirectRequest(tservers.TestMaster): f.request.port = self.redirect_port super(MasterRedirectRequest, self).handle_request(f) + @controller.handler def handle_response(self, f): f.response.content = str(f.client_conn.address.port) f.response.headers["server-conn-id"] = str(f.server_conn.source_address.port) @@ -689,10 +692,9 @@ class MasterStreamRequest(tservers.TestMaster): """ Enables the stream flag on the flow for all requests """ - + @controller.handler def handle_responseheaders(self, f): f.response.stream = True - f.reply() class TestStreamRequest(tservers.HTTPProxyTest): @@ -739,7 +741,7 @@ class TestStreamRequest(tservers.HTTPProxyTest): class MasterFakeResponse(tservers.TestMaster): - + @controller.handler def handle_request(self, f): resp = HTTPResponse.wrap(netlib.tutils.tresp()) f.reply(resp) @@ -767,6 +769,7 @@ class TestServerConnect(tservers.HTTPProxyTest): class MasterKillRequest(tservers.TestMaster): + @controller.handler def handle_request(self, f): f.reply(Kill) @@ -783,6 +786,7 @@ class TestKillRequest(tservers.HTTPProxyTest): class MasterKillResponse(tservers.TestMaster): + @controller.handler def handle_response(self, f): f.reply(Kill) @@ -812,6 +816,7 @@ class TestTransparentResolveError(tservers.TransparentProxyTest): class MasterIncomplete(tservers.TestMaster): + @controller.handler def handle_request(self, f): resp = HTTPResponse.wrap(netlib.tutils.tresp()) resp.content = None @@ -930,7 +935,9 @@ class TestProxyChainingSSLReconnect(tservers.HTTPUpstreamProxyTest): k = [0] # variable scope workaround: put into array _func = getattr(master, attr) - def handler(f): + @controller.handler + def handler(*args): + f = args[-1] k[0] += 1 if not (k[0] in exclude): f.client_conn.finish() @@ -940,11 +947,14 @@ class TestProxyChainingSSLReconnect(tservers.HTTPUpstreamProxyTest): setattr(master, attr, handler) - kill_requests(self.chain[1].tmaster, "handle_request", - exclude=[ - # fail first request - 2, # allow second request - ]) + kill_requests( + self.chain[1].tmaster, + "handle_request", + exclude = [ + # fail first request + 2, # allow second request + ] + ) kill_requests(self.chain[0].tmaster, "handle_request", exclude=[ @@ -1004,10 +1014,10 @@ class AddUpstreamCertsToClientChainMixin: ssl = True servercert = tutils.test_data.path("data/trusted-server.crt") ssloptions = pathod.SSLOptions( - cn="trusted-cert", - certs=[ - ("trusted-cert", servercert) - ] + cn="trusted-cert", + certs=[ + ("trusted-cert", servercert) + ] ) def test_add_upstream_certs_to_client_chain(self): diff --git a/test/mitmproxy/tservers.py b/test/mitmproxy/tservers.py index c9d68cfd..51f4109d 100644 --- a/test/mitmproxy/tservers.py +++ b/test/mitmproxy/tservers.py @@ -39,13 +39,13 @@ class TestMaster(flow.FlowMaster): self.apps.add(errapp, "errapp", 80) self.clear_log() + @controller.handler def handle_request(self, f): flow.FlowMaster.handle_request(self, f) - f.reply() + @controller.handler def handle_response(self, f): flow.FlowMaster.handle_response(self, f) - f.reply() def clear_log(self): self.log = [] |