diff options
Diffstat (limited to 'mitmproxy/test/taddons.py')
| -rw-r--r-- | mitmproxy/test/taddons.py | 41 |
1 files changed, 11 insertions, 30 deletions
diff --git a/mitmproxy/test/taddons.py b/mitmproxy/test/taddons.py index 82a935d2..509f8d53 100644 --- a/mitmproxy/test/taddons.py +++ b/mitmproxy/test/taddons.py @@ -1,4 +1,5 @@ import contextlib +import asyncio import sys import mitmproxy.master @@ -16,10 +17,6 @@ class TestAddons(addonmanager.AddonManager): def trigger(self, event, *args, **kwargs): if event == "log": self.master.logs.append(args[0]) - elif event == "tick" and not args and not kwargs: - pass - else: - self.master.events.append((event, args, kwargs)) super().trigger(event, *args, **kwargs) @@ -27,7 +24,6 @@ class RecordingMaster(mitmproxy.master.Master): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.addons = TestAddons(self) - self.events = [] self.logs = [] def dump_log(self, outf=sys.stdout): @@ -42,10 +38,12 @@ class RecordingMaster(mitmproxy.master.Master): return True return False - def has_event(self, name): - for i in self.events: - if i[0] == name: + async def await_log(self, txt, level=None): + for i in range(20): + if self.has_log(txt, level): return True + else: + await asyncio.sleep(0.1) return False def clear(self): @@ -65,7 +63,6 @@ class context: options ) self.options = self.master.options - self.wrapped = None if loadcore: self.master.addons.add(core.Core()) @@ -73,20 +70,10 @@ class context: for a in addons: self.master.addons.add(a) - def ctx(self): - """ - Returns a new handler context. - """ - return self.master.handlecontext() - def __enter__(self): - self.wrapped = self.ctx() - self.wrapped.__enter__() return self def __exit__(self, exc_type, exc_value, traceback): - self.wrapped.__exit__(exc_type, exc_value, traceback) - self.wrapped = None return False @contextlib.contextmanager @@ -114,22 +101,16 @@ class context: if addon not in self.master.addons: self.master.addons.register(addon) with self.options.rollback(kwargs.keys(), reraise=True): - self.options.update(**kwargs) - self.master.addons.invoke_addon( - addon, - "configure", - kwargs.keys() - ) + if kwargs: + self.options.update(**kwargs) + else: + self.master.addons.invoke_addon(addon, "configure", {}) def script(self, path): """ Loads a script from path, and returns the enclosed addon. """ - sc = script.Script(path) - loader = addonmanager.Loader(self.master) - self.master.addons.invoke_addon(sc, "load", loader) - self.configure(sc) - self.master.addons.invoke_addon(sc, "tick") + sc = script.Script(path, False) return sc.addons[0] if sc.addons else None def invoke(self, addon, event, *args, **kwargs): |