diff options
Diffstat (limited to 'test/mitmproxy/mastertest.py')
| -rw-r--r-- | test/mitmproxy/mastertest.py | 31 |
1 files changed, 23 insertions, 8 deletions
diff --git a/test/mitmproxy/mastertest.py b/test/mitmproxy/mastertest.py index 9e726a32..d1fe8cb4 100644 --- a/test/mitmproxy/mastertest.py +++ b/test/mitmproxy/mastertest.py @@ -3,24 +3,30 @@ import mock from . import tutils import netlib.tutils -from mitmproxy import flow, proxy, models +from mitmproxy.flow import master +from mitmproxy import flow, proxy, models, controller class MasterTest: + def invoke(self, master, handler, message): + with master.handlecontext(): + func = getattr(master, handler) + func(message) + message.reply = controller.DummyReply() + def cycle(self, master, content): f = tutils.tflow(req=netlib.tutils.treq(content=content)) l = proxy.Log("connect") l.reply = mock.MagicMock() master.log(l) - master.clientconnect(f.client_conn) - master.serverconnect(f.server_conn) - master.request(f) + self.invoke(master, "clientconnect", f.client_conn) + self.invoke(master, "clientconnect", f.client_conn) + self.invoke(master, "serverconnect", f.server_conn) + self.invoke(master, "request", f) if not f.error: f.response = models.HTTPResponse.wrap(netlib.tutils.tresp(content=content)) - f.reply.acked = False - f = master.response(f) - f.client_conn.reply.acked = False - master.clientdisconnect(f.client_conn) + self.invoke(master, "response", f) + self.invoke(master, "clientdisconnect", f) return f def dummy_cycle(self, master, n, content): @@ -34,3 +40,12 @@ class MasterTest: t = tutils.tflow(resp=True) fw.add(t) f.close() + + +class RecordingMaster(master.FlowMaster): + def __init__(self, *args, **kwargs): + master.FlowMaster.__init__(self, *args, **kwargs) + self.event_log = [] + + def add_log(self, e, level): + self.event_log.append((level, e)) |