diff options
Diffstat (limited to 'mitmproxy/test/taddons.py')
| -rw-r--r-- | mitmproxy/test/taddons.py | 34 |
1 files changed, 20 insertions, 14 deletions
diff --git a/mitmproxy/test/taddons.py b/mitmproxy/test/taddons.py index c3b65e92..8bc174c7 100644 --- a/mitmproxy/test/taddons.py +++ b/mitmproxy/test/taddons.py @@ -3,36 +3,42 @@ import contextlib import mitmproxy.master import mitmproxy.options from mitmproxy import proxy +from mitmproxy import addonmanager from mitmproxy import eventsequence -class _AddonWrapper: - def __init__(self, master, addons): - self.master = master - self.addons = addons +class TestAddons(addonmanager.AddonManager): + def __init__(self, master): + super().__init__(master) def trigger(self, event, *args, **kwargs): - self.master.events.append((event, args, kwargs)) - return self.addons.trigger(event, *args, **kwargs) - - def __getattr__(self, attr): - return getattr(self.addons, attr) + if event == "log": + self.master.logs.append(args[0]) + else: + self.master.events.append((event, args, kwargs)) + super().trigger(event, *args, **kwargs) class RecordingMaster(mitmproxy.master.Master): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - self.addons = _AddonWrapper(self, self.addons) + self.addons = TestAddons(self) self.events = [] self.logs = [] + def has_log(self, txt, level=None): + for i in self.logs: + if level and i.level != level: + continue + if txt.lower() in i.msg.lower(): + return True + return False + def has_event(self, name): for i in self.events: if i[0] == name: return True - - def add_log(self, e, level): - self.logs.append((level, e)) + return False def clear(self): self.logs = [] @@ -67,7 +73,7 @@ class context: Cycles the flow through the events for the flow. Stops if a reply is taken (as in flow interception). """ - f.reply._state = "handled" + f.reply._state = "start" for evt, arg in eventsequence.iterate(f): h = getattr(addon, evt, None) if h: |