diff options
| -rw-r--r-- | mitmproxy/addonmanager.py | 30 | ||||
| -rw-r--r-- | mitmproxy/command.py | 6 | ||||
| -rw-r--r-- | mitmproxy/master.py | 21 | ||||
| -rw-r--r-- | mitmproxy/test/taddons.py | 11 | ||||
| -rw-r--r-- | setup.py | 1 | ||||
| -rw-r--r-- | test/mitmproxy/addons/test_script.py | 5 | ||||
| -rw-r--r-- | test/mitmproxy/proxy/protocol/test_websocket.py | 1 |
7 files changed, 21 insertions, 54 deletions
diff --git a/mitmproxy/addonmanager.py b/mitmproxy/addonmanager.py index 7c446c80..9d0e7069 100644 --- a/mitmproxy/addonmanager.py +++ b/mitmproxy/addonmanager.py @@ -8,7 +8,6 @@ from mitmproxy import exceptions from mitmproxy import eventsequence from mitmproxy import controller from mitmproxy import flow -from mitmproxy import log from . import ctx import pprint @@ -38,9 +37,6 @@ def cut_traceback(tb, func_name): class StreamLog: - """ - A class for redirecting output using contextlib. - """ def __init__(self, lg): self.log = lg @@ -183,9 +179,8 @@ class AddonManager: Add addons to the end of the chain, and run their load event. If any addon has sub-addons, they are registered. """ - with self.master.handlecontext(): - for i in addons: - self.chain.append(self.register(i)) + for i in addons: + self.chain.append(self.register(i)) def remove(self, addon): """ @@ -201,8 +196,7 @@ class AddonManager: raise exceptions.AddonManagerError("No such addon: %s" % n) self.chain = [i for i in self.chain if i is not a] del self.lookup[_get_name(a)] - with self.master.handlecontext(): - self.invoke_addon(a, "done") + self.invoke_addon(a, "done") def __len__(self): return len(self.chain) @@ -245,8 +239,7 @@ class AddonManager: def invoke_addon(self, addon, name, *args, **kwargs): """ - Invoke an event on an addon and all its children. This method must - run within an established handler context. + Invoke an event on an addon and all its children. """ if name not in eventsequence.Events: name = "event_" + name @@ -268,12 +261,11 @@ class AddonManager: def trigger(self, name, *args, **kwargs): """ - Establish a handler context and trigger an event across all addons + Trigger an event across all addons. """ - with self.master.handlecontext(): - for i in self.chain: - try: - with safecall(): - self.invoke_addon(i, name, *args, **kwargs) - except exceptions.AddonHalt: - return + for i in self.chain: + try: + with safecall(): + self.invoke_addon(i, name, *args, **kwargs) + except exceptions.AddonHalt: + return diff --git a/mitmproxy/command.py b/mitmproxy/command.py index 114e882d..35a7f42a 100644 --- a/mitmproxy/command.py +++ b/mitmproxy/command.py @@ -95,11 +95,7 @@ class Command: Call the command with a list of arguments. At this point, all arguments are strings. """ - pargs = self.prepare_args(args) - - with self.manager.master.handlecontext(): - ret = self.func(*pargs) - + ret = self.func(*self.prepare_args(args)) if ret is None and self.returntype is None: return typ = mitmproxy.types.CommandTypes.get(self.returntype) diff --git a/mitmproxy/master.py b/mitmproxy/master.py index 323f73dc..4e17e538 100644 --- a/mitmproxy/master.py +++ b/mitmproxy/master.py @@ -1,5 +1,4 @@ import threading -import contextlib import asyncio import logging @@ -58,6 +57,10 @@ class Master: self.waiting_flows = [] self.log = log.Log(self) + mitmproxy_ctx.master = self + mitmproxy_ctx.log = self.log + mitmproxy_ctx.options = self.options + @property def server(self): return self._server @@ -67,22 +70,6 @@ class Master: server.set_channel(self.channel) self._server = server - @contextlib.contextmanager - def handlecontext(self): - # Handlecontexts also have to nest - leave cleanup to the outermost - if mitmproxy_ctx.master: - yield - return - mitmproxy_ctx.master = self - mitmproxy_ctx.log = self.log - mitmproxy_ctx.options = self.options - try: - yield - finally: - mitmproxy_ctx.master = None - mitmproxy_ctx.log = None - mitmproxy_ctx.options = None - def start(self): self.should_exit.clear() if self.server: diff --git a/mitmproxy/test/taddons.py b/mitmproxy/test/taddons.py index 6d439c4a..73e57456 100644 --- a/mitmproxy/test/taddons.py +++ b/mitmproxy/test/taddons.py @@ -74,7 +74,6 @@ class context: options ) self.options = self.master.options - self.wrapped = None if loadcore: self.master.addons.add(core.Core()) @@ -82,20 +81,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 @@ -88,6 +88,7 @@ setup( "flake8>=3.5, <3.6", "Flask>=0.10.1, <0.13", "mypy>=0.580,<0.581", + "pytest-asyncio>=0.8", "pytest-cov>=2.5.1,<3", "pytest-faulthandler>=1.3.1,<2", "pytest-timeout>=1.2.1,<2", diff --git a/test/mitmproxy/addons/test_script.py b/test/mitmproxy/addons/test_script.py index 17b1e372..3f0ce68c 100644 --- a/test/mitmproxy/addons/test_script.py +++ b/test/mitmproxy/addons/test_script.py @@ -180,11 +180,12 @@ class TestScriptLoader: 'recorder responseheaders', 'recorder response' ] - def test_script_run_nonexistent(self): + @pytest.mark.asyncio + async def test_script_run_nonexistent(self): sc = script.ScriptLoader() with taddons.context(sc) as tctx: sc.script_run([tflow.tflow(resp=True)], "/") - tctx.master.has_log("/: No such script") + assert await tctx.master.await_log("/: No such script") def test_simple(self): sc = script.ScriptLoader() diff --git a/test/mitmproxy/proxy/protocol/test_websocket.py b/test/mitmproxy/proxy/protocol/test_websocket.py index 347c0dfc..3ce1436a 100644 --- a/test/mitmproxy/proxy/protocol/test_websocket.py +++ b/test/mitmproxy/proxy/protocol/test_websocket.py @@ -47,6 +47,7 @@ class _WebSocketServerBase(net_tservers.ServerTestBase): class _WebSocketTestBase: + client = None @classmethod def setup_class(cls): |