diff options
Diffstat (limited to 'mitmproxy/test/taddons.py')
-rw-r--r-- | mitmproxy/test/taddons.py | 44 |
1 files changed, 36 insertions, 8 deletions
diff --git a/mitmproxy/test/taddons.py b/mitmproxy/test/taddons.py index 8bc174c7..3dbccba2 100644 --- a/mitmproxy/test/taddons.py +++ b/mitmproxy/test/taddons.py @@ -1,3 +1,4 @@ +import sys import contextlib import mitmproxy.master @@ -5,6 +6,7 @@ import mitmproxy.options from mitmproxy import proxy from mitmproxy import addonmanager from mitmproxy import eventsequence +from mitmproxy.addons import script class TestAddons(addonmanager.AddonManager): @@ -26,6 +28,10 @@ class RecordingMaster(mitmproxy.master.Master): self.events = [] self.logs = [] + def dump_log(self, outf=sys.stdout): + for i in self.logs: + print("%s: %s" % (i.level, i.msg), file=outf) + def has_log(self, txt, level=None): for i in self.logs: if level and i.level != level: @@ -51,14 +57,21 @@ class context: provides a number of helper methods for common testing scenarios. """ def __init__(self, master = None, options = None): - self.options = options or mitmproxy.options.Options() + options = options or mitmproxy.options.Options() self.master = master or RecordingMaster( options, proxy.DummyServer(options) ) + self.options = self.master.options self.wrapped = None + def ctx(self): + """ + Returns a new handler context. + """ + return self.master.handlecontext() + def __enter__(self): - self.wrapped = self.master.handlecontext() + self.wrapped = self.ctx() self.wrapped.__enter__() return self @@ -75,11 +88,13 @@ class context: """ f.reply._state = "start" for evt, arg in eventsequence.iterate(f): - h = getattr(addon, evt, None) - if h: - h(arg) - if f.reply.state == "taken": - return + self.master.addons.invoke_addon( + addon, + evt, + arg + ) + if f.reply.state == "taken": + return def configure(self, addon, **kwargs): """ @@ -89,4 +104,17 @@ class context: """ with self.options.rollback(kwargs.keys(), reraise=True): self.options.update(**kwargs) - addon.configure(self.options, kwargs.keys()) + self.master.addons.invoke_addon( + addon, + "configure", + self.options, + kwargs.keys() + ) + + def script(self, path): + sc = script.Script(path) + loader = addonmanager.Loader(self.master) + sc.load(loader) + for a in addonmanager.traverse(sc.addons): + getattr(a, "load", lambda x: None)(loader) + return sc |