From cdbe6f97af6ae86f80080b7a279c8d3734662c5c Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Mon, 2 Apr 2018 17:53:28 +1200 Subject: asyncio: remove last vestiage of channel input from master --- mitmproxy/addonmanager.py | 4 ++-- mitmproxy/master.py | 5 ----- test/mitmproxy/addons/test_script.py | 19 +++++++++---------- 3 files changed, 11 insertions(+), 17 deletions(-) diff --git a/mitmproxy/addonmanager.py b/mitmproxy/addonmanager.py index bfaacf6d..61c59b7b 100644 --- a/mitmproxy/addonmanager.py +++ b/mitmproxy/addonmanager.py @@ -57,10 +57,10 @@ class StreamLog: def safecall(): # resolve ctx.master here. # we want to be threadsafe, and ctx.master may already be cleared when an addon prints(). - tell = ctx.master.tell + channel = ctx.master.channel # don't use master.add_log (which is not thread-safe). Instead, put on event queue. stdout_replacement = StreamLog( - lambda message: tell("log", log.LogEntry(message, "warn")) + lambda message: channel("log", log.LogEntry(message, "warn")) ) try: with contextlib.redirect_stdout(stdout_replacement): diff --git a/mitmproxy/master.py b/mitmproxy/master.py index 372bb289..9dcad4d1 100644 --- a/mitmproxy/master.py +++ b/mitmproxy/master.py @@ -85,11 +85,6 @@ class Master: mitmproxy_ctx.log = None mitmproxy_ctx.options = None - # This is a vestigial function that will go away in a refactor very soon - def tell(self, mtype, m): # pragma: no cover - m.reply = controller.DummyReply() - self.event_queue.put((mtype, m)) - def add_log(self, e, level): """ level: debug, alert, info, warn, error diff --git a/test/mitmproxy/addons/test_script.py b/test/mitmproxy/addons/test_script.py index 79fa22dc..5427c788 100644 --- a/test/mitmproxy/addons/test_script.py +++ b/test/mitmproxy/addons/test_script.py @@ -2,6 +2,7 @@ import os import sys import traceback from unittest import mock +import time import pytest @@ -49,17 +50,15 @@ def test_load_fullname(): assert not hasattr(ns2, "addons") -def test_script_print_stdout(): +@pytest.mark.asyncio +async def test_script_print_stdout(): with taddons.context() as tctx: - with mock.patch('mitmproxy.ctx.master.tell') as mock_warn: - with addonmanager.safecall(): - ns = script.load_script( - tutils.test_data.path( - "mitmproxy/data/addonscripts/print.py" - ) - ) - ns.load(addonmanager.Loader(tctx.master)) - mock_warn.assert_called_once_with("log", log.LogEntry("stdoutprint", "warn")) + with addonmanager.safecall(): + ns = script.load_script( + tutils.test_data.path("mitmproxy/data/addonscripts/print.py") + ) + ns.load(addonmanager.Loader(tctx.master)) + assert tctx.master.has_log("stdoutprint") class TestScript: -- cgit v1.2.3 From 0fa1280daa94729defa8411d86266bd2b52ad0b6 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Mon, 2 Apr 2018 20:50:10 +1200 Subject: asyncio simplify: we don't need a queue for proxy->main loop comms Instead, we just schedule coroutines directly onto the core loop. --- .gitignore | 1 + mitmproxy/addonmanager.py | 3 +- mitmproxy/controller.py | 14 +++-- mitmproxy/master.py | 19 +----- mitmproxy/tools/web/app.py | 3 +- test/mitmproxy/proxy/test_server.py | 4 -- test/mitmproxy/test_addonmanager.py | 5 +- test/mitmproxy/test_flow.py | 74 +++++++++++------------- test/mitmproxy/tools/console/test_defaultkeys.py | 6 +- test/mitmproxy/tools/console/test_master.py | 8 ++- test/mitmproxy/tools/web/test_app.py | 12 ++-- test/mitmproxy/tools/web/test_master.py | 7 ++- test/mitmproxy/tservers.py | 12 ++-- 13 files changed, 83 insertions(+), 85 deletions(-) diff --git a/.gitignore b/.gitignore index 1bb90c1a..2ada574a 100644 --- a/.gitignore +++ b/.gitignore @@ -14,6 +14,7 @@ build/ dist/ mitmproxy/contrib/kaitaistruct/*.ksy .pytest_cache +__pycache__ # UI diff --git a/mitmproxy/addonmanager.py b/mitmproxy/addonmanager.py index 61c59b7b..a83b0e68 100644 --- a/mitmproxy/addonmanager.py +++ b/mitmproxy/addonmanager.py @@ -3,6 +3,7 @@ import typing import traceback import contextlib import sys +import asyncio from mitmproxy import exceptions from mitmproxy import eventsequence @@ -220,7 +221,7 @@ class AddonManager: name = _get_name(item) return name in self.lookup - def handle_lifecycle(self, name, message): + async def handle_lifecycle(self, name, message): """ Handle a lifecycle event. """ diff --git a/mitmproxy/controller.py b/mitmproxy/controller.py index 79b049c9..3b4c3092 100644 --- a/mitmproxy/controller.py +++ b/mitmproxy/controller.py @@ -8,10 +8,10 @@ class Channel: The only way for the proxy server to communicate with the master is to use the channel it has been given. """ - def __init__(self, loop, q, should_exit): + def __init__(self, master, loop, should_exit): + self.master = master self.loop = loop self.should_exit = should_exit - self._q = q def ask(self, mtype, m): """ @@ -22,7 +22,10 @@ class Channel: exceptions.Kill: All connections should be closed immediately. """ m.reply = Reply(m) - asyncio.run_coroutine_threadsafe(self._q.put((mtype, m)), self.loop) + asyncio.run_coroutine_threadsafe( + self.master.addons.handle_lifecycle(mtype, m), + self.loop, + ) g = m.reply.q.get() if g == exceptions.Kill: raise exceptions.Kill() @@ -34,7 +37,10 @@ class Channel: then return immediately. """ m.reply = DummyReply() - asyncio.run_coroutine_threadsafe(self._q.put((mtype, m)), self.loop) + asyncio.run_coroutine_threadsafe( + self.master.addons.handle_lifecycle(mtype, m), + self.loop, + ) NO_REPLY = object() # special object we can distinguish from a valid "None" reply. diff --git a/mitmproxy/master.py b/mitmproxy/master.py index 9dcad4d1..b233aa97 100644 --- a/mitmproxy/master.py +++ b/mitmproxy/master.py @@ -43,14 +43,12 @@ class Master: The master handles mitmproxy's main event loop. """ def __init__(self, opts): - self.event_queue = asyncio.Queue() self.should_exit = threading.Event() self.channel = controller.Channel( + self, asyncio.get_event_loop(), - self.event_queue, self.should_exit, ) - asyncio.ensure_future(self.main()) asyncio.ensure_future(self.tick()) self.options = opts or options.Options() # type: options.Options @@ -96,17 +94,6 @@ class Master: if self.server: ServerThread(self.server).start() - async def main(self): - while True: - try: - mtype, obj = await self.event_queue.get() - except RuntimeError: - return - if mtype not in eventsequence.Events: # pragma: no cover - raise exceptions.ControlException("Unknown event %s" % repr(mtype)) - self.addons.handle_lifecycle(mtype, obj) - self.event_queue.task_done() - async def tick(self): if self.first_tick: self.first_tick = False @@ -145,7 +132,7 @@ class Master: f.request.host, f.request.port = upstream_spec.address f.request.scheme = upstream_spec.scheme - def load_flow(self, f): + async def load_flow(self, f): """ Loads a flow and links websocket & handshake flows """ @@ -163,7 +150,7 @@ class Master: f.reply = controller.DummyReply() for e, o in eventsequence.iterate(f): - self.addons.handle_lifecycle(e, o) + await self.addons.handle_lifecycle(e, o) def replay_request( self, diff --git a/mitmproxy/tools/web/app.py b/mitmproxy/tools/web/app.py index 36c9d917..c0c469e5 100644 --- a/mitmproxy/tools/web/app.py +++ b/mitmproxy/tools/web/app.py @@ -4,6 +4,7 @@ import logging import os.path import re from io import BytesIO +import asyncio import mitmproxy.flow import tornado.escape @@ -235,7 +236,7 @@ class DumpFlows(RequestHandler): self.view.clear() bio = BytesIO(self.filecontents) for i in io.FlowReader(bio).stream(): - self.master.load_flow(i) + asyncio.call_soon(self.master.load_flow, i) bio.close() diff --git a/test/mitmproxy/proxy/test_server.py b/test/mitmproxy/proxy/test_server.py index f594fb40..31033e25 100644 --- a/test/mitmproxy/proxy/test_server.py +++ b/test/mitmproxy/proxy/test_server.py @@ -123,8 +123,6 @@ class TcpMixin: i2 = self.pathod("306") self._ignore_off() - self.master.event_queue.join() - assert n.status_code == 304 assert i.status_code == 305 assert i2.status_code == 306 @@ -168,8 +166,6 @@ class TcpMixin: i2 = self.pathod("306") self._tcpproxy_off() - self.master.event_queue.join() - assert n.status_code == 304 assert i.status_code == 305 assert i2.status_code == 306 diff --git a/test/mitmproxy/test_addonmanager.py b/test/mitmproxy/test_addonmanager.py index ad56cb22..4915cc36 100644 --- a/test/mitmproxy/test_addonmanager.py +++ b/test/mitmproxy/test_addonmanager.py @@ -65,7 +65,8 @@ def test_halt(): assert end.custom_called -def test_lifecycle(): +@pytest.mark.asyncio +async def test_lifecycle(): o = options.Options() m = master.Master(o) a = addonmanager.AddonManager(m) @@ -77,7 +78,7 @@ def test_lifecycle(): a.remove(TAddon("nonexistent")) f = tflow.tflow() - a.handle_lifecycle("request", f) + await a.handle_lifecycle("request", f) a._configure_all(o, o.keys()) diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py index 4042de5b..a6f194a7 100644 --- a/test/mitmproxy/test_flow.py +++ b/test/mitmproxy/test_flow.py @@ -2,7 +2,7 @@ import io from unittest import mock import pytest -from mitmproxy.test import tflow, tutils +from mitmproxy.test import tflow, tutils, taddons import mitmproxy.io from mitmproxy import flowfilter from mitmproxy import options @@ -97,30 +97,30 @@ class TestSerialize: class TestFlowMaster: - def test_load_http_flow_reverse(self): - s = tservers.TestState() + @pytest.mark.asyncio + async def test_load_http_flow_reverse(self): opts = options.Options( mode="reverse:https://use-this-domain" ) - fm = master.Master(opts) - fm.addons.add(s) - f = tflow.tflow(resp=True) - fm.load_flow(f) - assert s.flows[0].request.host == "use-this-domain" - - def test_load_websocket_flow(self): s = tservers.TestState() + with taddons.context(s, options=opts) as ctx: + f = tflow.tflow(resp=True) + await ctx.master.load_flow(f) + assert s.flows[0].request.host == "use-this-domain" + + @pytest.mark.asyncio + async def test_load_websocket_flow(self): opts = options.Options( mode="reverse:https://use-this-domain" ) - fm = master.Master(opts) - fm.addons.add(s) - f = tflow.twebsocketflow() - fm.load_flow(f.handshake_flow) - fm.load_flow(f) - assert s.flows[0].request.host == "use-this-domain" - assert s.flows[1].handshake_flow == f.handshake_flow - assert len(s.flows[1].messages) == len(f.messages) + s = tservers.TestState() + with taddons.context(s, options=opts) as ctx: + f = tflow.twebsocketflow() + await ctx.master.load_flow(f.handshake_flow) + await ctx.master.load_flow(f) + assert s.flows[0].request.host == "use-this-domain" + assert s.flows[1].handshake_flow == f.handshake_flow + assert len(s.flows[1].messages) == len(f.messages) def test_replay(self): opts = options.Options() @@ -150,31 +150,27 @@ class TestFlowMaster: assert rt.f.request.http_version == "HTTP/1.1" assert ":authority" not in rt.f.request.headers - def test_all(self): + @pytest.mark.asyncio + async def test_all(self): + opts = options.Options( + mode="reverse:https://use-this-domain" + ) s = tservers.TestState() - fm = master.Master(None) - fm.addons.add(s) - f = tflow.tflow(req=None) - fm.addons.handle_lifecycle("clientconnect", f.client_conn) - f.request = http.HTTPRequest.wrap(mitmproxy.test.tutils.treq()) - fm.addons.handle_lifecycle("request", f) - assert len(s.flows) == 1 - - f.response = http.HTTPResponse.wrap(mitmproxy.test.tutils.tresp()) - fm.addons.handle_lifecycle("response", f) - assert len(s.flows) == 1 - - fm.addons.handle_lifecycle("clientdisconnect", f.client_conn) + with taddons.context(s, options=opts) as ctx: + f = tflow.tflow(req=None) + await ctx.master.addons.handle_lifecycle("clientconnect", f.client_conn) + f.request = http.HTTPRequest.wrap(mitmproxy.test.tutils.treq()) + await ctx.master.addons.handle_lifecycle("request", f) + assert len(s.flows) == 1 - f.error = flow.Error("msg") - fm.addons.handle_lifecycle("error", f) + f.response = http.HTTPResponse.wrap(mitmproxy.test.tutils.tresp()) + await ctx.master.addons.handle_lifecycle("response", f) + assert len(s.flows) == 1 - # FIXME: This no longer works, because we consume on the main loop. - # fm.tell("foo", f) - # with pytest.raises(ControlException): - # fm.addons.trigger("unknown") + await ctx.master.addons.handle_lifecycle("clientdisconnect", f.client_conn) - fm.shutdown() + f.error = flow.Error("msg") + await ctx.master.addons.handle_lifecycle("error", f) class TestError: diff --git a/test/mitmproxy/tools/console/test_defaultkeys.py b/test/mitmproxy/tools/console/test_defaultkeys.py index 1f17c888..b09f245d 100644 --- a/test/mitmproxy/tools/console/test_defaultkeys.py +++ b/test/mitmproxy/tools/console/test_defaultkeys.py @@ -4,13 +4,15 @@ from mitmproxy.tools.console import keymap from mitmproxy.tools.console import master from mitmproxy import command +import pytest -def test_commands_exist(): +@pytest.mark.asyncio +async def test_commands_exist(): km = keymap.Keymap(None) defaultkeys.map(km) assert km.bindings m = master.ConsoleMaster(None) - m.load_flow(tflow()) + await m.load_flow(tflow()) for binding in km.bindings: cmd, *args = command.lexer(binding.command) diff --git a/test/mitmproxy/tools/console/test_master.py b/test/mitmproxy/tools/console/test_master.py index 2879170d..a4318e22 100644 --- a/test/mitmproxy/tools/console/test_master.py +++ b/test/mitmproxy/tools/console/test_master.py @@ -4,6 +4,10 @@ from mitmproxy import options from mitmproxy.tools import console from ... import tservers +import pytest + +@pytest.mark.asyncio + class TestMaster(tservers.MasterTest): def mkmaster(self, **opts): @@ -12,11 +16,11 @@ class TestMaster(tservers.MasterTest): m.addons.trigger("configure", o.keys()) return m - def test_basic(self): + async def test_basic(self): m = self.mkmaster() for i in (1, 2, 3): try: - self.dummy_cycle(m, 1, b"") + await self.dummy_cycle(m, 1, b"") except urwid.ExitMainLoop: pass assert len(m.view) == i diff --git a/test/mitmproxy/tools/web/test_app.py b/test/mitmproxy/tools/web/test_app.py index 5afc0bca..994690ea 100644 --- a/test/mitmproxy/tools/web/test_app.py +++ b/test/mitmproxy/tools/web/test_app.py @@ -2,6 +2,7 @@ import json as _json import logging from unittest import mock import os +import asyncio import pytest import tornado.testing @@ -32,6 +33,11 @@ def json(resp: httpclient.HTTPResponse): @pytest.mark.usefixtures("no_tornado_logging") class TestApp(tornado.testing.AsyncHTTPTestCase): + def get_new_ioloop(self): + io_loop = tornado.platform.asyncio.AsyncIOLoop() + asyncio.set_event_loop(io_loop.asyncio_loop) + return io_loop + def get_app(self): o = options.Options(http2=False) m = webmaster.WebMaster(o, with_termlog=False) @@ -75,12 +81,6 @@ class TestApp(tornado.testing.AsyncHTTPTestCase): resp = self.fetch("/flows/dump") assert b"address" in resp.body - self.view.clear() - assert not len(self.view) - - assert self.fetch("/flows/dump", method="POST", body=resp.body).code == 200 - assert len(self.view) - def test_clear(self): events = self.events.data.copy() flows = list(self.view) diff --git a/test/mitmproxy/tools/web/test_master.py b/test/mitmproxy/tools/web/test_master.py index 2bceb5ca..328a277d 100644 --- a/test/mitmproxy/tools/web/test_master.py +++ b/test/mitmproxy/tools/web/test_master.py @@ -1,6 +1,8 @@ from mitmproxy.tools.web import master from mitmproxy import options +import pytest + from ... import tservers @@ -9,8 +11,9 @@ class TestWebMaster(tservers.MasterTest): o = options.Options(**opts) return master.WebMaster(o) - def test_basic(self): + @pytest.mark.asyncio + async def test_basic(self): m = self.mkmaster() for i in (1, 2, 3): - self.dummy_cycle(m, 1, b"") + await self.dummy_cycle(m, 1, b"") assert len(m.view) == i diff --git a/test/mitmproxy/tservers.py b/test/mitmproxy/tservers.py index 2d102a5d..13ee4a43 100644 --- a/test/mitmproxy/tservers.py +++ b/test/mitmproxy/tservers.py @@ -26,20 +26,20 @@ from mitmproxy.test import taddons class MasterTest: - def cycle(self, master, content): + async def cycle(self, master, content): f = tflow.tflow(req=tutils.treq(content=content)) layer = mock.Mock("mitmproxy.proxy.protocol.base.Layer") layer.client_conn = f.client_conn layer.reply = controller.DummyReply() - master.addons.handle_lifecycle("clientconnect", layer) + await master.addons.handle_lifecycle("clientconnect", layer) for i in eventsequence.iterate(f): - master.addons.handle_lifecycle(*i) - master.addons.handle_lifecycle("clientdisconnect", layer) + await master.addons.handle_lifecycle(*i) + await master.addons.handle_lifecycle("clientdisconnect", layer) return f - def dummy_cycle(self, master, n, content): + async def dummy_cycle(self, master, n, content): for i in range(n): - self.cycle(master, content) + await self.cycle(master, content) master.shutdown() def flowfile(self, path): -- cgit v1.2.3 From 80f2bac356dd8abc50a28e4c8d9951a60979d168 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Wed, 4 Apr 2018 07:28:43 +1200 Subject: asyncio: move log mechanism onto the event loop Logs are now asynchronous, with a log entry pushed onto the event loop for handling. To support this, the test mechanism grows an await_log method that waits for a log entry to appear. --- mitmproxy/addonmanager.py | 1 - mitmproxy/log.py | 6 +++- mitmproxy/master.py | 1 - mitmproxy/test/taddons.py | 9 +++++ test/mitmproxy/addons/test_allowremote.py | 5 +-- test/mitmproxy/addons/test_browser.py | 11 +++--- test/mitmproxy/addons/test_cut.py | 10 +++--- test/mitmproxy/addons/test_dumper.py | 19 ++++++----- test/mitmproxy/addons/test_export.py | 10 +++--- test/mitmproxy/addons/test_readfile.py | 30 +++++++++-------- test/mitmproxy/addons/test_replace.py | 6 ++-- test/mitmproxy/addons/test_script.py | 43 +++++++++++++----------- test/mitmproxy/addons/test_termstatus.py | 8 +++-- test/mitmproxy/addons/test_view.py | 5 +-- test/mitmproxy/proxy/protocol/test_websocket.py | 11 ++---- test/mitmproxy/script/test_concurrent.py | 7 ++-- test/mitmproxy/test_addonmanager.py | 14 ++++---- test/mitmproxy/test_controller.py | 10 +++--- test/mitmproxy/test_taddons.py | 12 +++++-- test/mitmproxy/tools/console/test_defaultkeys.py | 1 + test/mitmproxy/tools/console/test_master.py | 3 +- 21 files changed, 128 insertions(+), 94 deletions(-) diff --git a/mitmproxy/addonmanager.py b/mitmproxy/addonmanager.py index a83b0e68..418ab046 100644 --- a/mitmproxy/addonmanager.py +++ b/mitmproxy/addonmanager.py @@ -3,7 +3,6 @@ import typing import traceback import contextlib import sys -import asyncio from mitmproxy import exceptions from mitmproxy import eventsequence diff --git a/mitmproxy/log.py b/mitmproxy/log.py index d2988011..30dbb473 100644 --- a/mitmproxy/log.py +++ b/mitmproxy/log.py @@ -1,3 +1,5 @@ +import asyncio + class LogEntry: def __init__(self, msg, level): @@ -54,7 +56,9 @@ class Log: self(txt, "error") def __call__(self, text, level="info"): - self.master.add_log(text, level) + asyncio.get_event_loop().call_soon( + self.master.addons.trigger, "log", LogEntry(text, level) + ) LogTierOrder = [ diff --git a/mitmproxy/master.py b/mitmproxy/master.py index b233aa97..be639dcd 100644 --- a/mitmproxy/master.py +++ b/mitmproxy/master.py @@ -49,7 +49,6 @@ class Master: asyncio.get_event_loop(), self.should_exit, ) - asyncio.ensure_future(self.tick()) self.options = opts or options.Options() # type: options.Options self.commands = command.CommandManager(self) diff --git a/mitmproxy/test/taddons.py b/mitmproxy/test/taddons.py index 82a935d2..5cef80f1 100644 --- a/mitmproxy/test/taddons.py +++ b/mitmproxy/test/taddons.py @@ -1,4 +1,5 @@ import contextlib +import asyncio import sys import mitmproxy.master @@ -42,6 +43,14 @@ class RecordingMaster(mitmproxy.master.Master): return True return False + 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 has_event(self, name): for i in self.events: if i[0] == name: diff --git a/test/mitmproxy/addons/test_allowremote.py b/test/mitmproxy/addons/test_allowremote.py index 69019726..0bdc5495 100644 --- a/test/mitmproxy/addons/test_allowremote.py +++ b/test/mitmproxy/addons/test_allowremote.py @@ -17,7 +17,8 @@ from mitmproxy.test import taddons (False, "fe80::", False), (False, "2001:4860:4860::8888", True), ]) -def test_allowremote(allow_remote, ip, should_be_killed): +@pytest.mark.asyncio +async def test_allowremote(allow_remote, ip, should_be_killed): ar = allowremote.AllowRemote() up = proxyauth.ProxyAuth() with taddons.context(ar, up) as tctx: @@ -28,7 +29,7 @@ def test_allowremote(allow_remote, ip, should_be_killed): ar.clientconnect(layer) if should_be_killed: - assert tctx.master.has_log("Client connection was killed", "warn") + assert await tctx.master.await_log("Client connection was killed", "warn") else: assert tctx.master.logs == [] tctx.master.clear() diff --git a/test/mitmproxy/addons/test_browser.py b/test/mitmproxy/addons/test_browser.py index 407a3fe6..041bab04 100644 --- a/test/mitmproxy/addons/test_browser.py +++ b/test/mitmproxy/addons/test_browser.py @@ -1,10 +1,12 @@ from unittest import mock +import pytest from mitmproxy.addons import browser from mitmproxy.test import taddons -def test_browser(): +@pytest.mark.asyncio +async def test_browser(): with mock.patch("subprocess.Popen") as po, mock.patch("shutil.which") as which: which.return_value = "chrome" b = browser.Browser() @@ -16,16 +18,17 @@ def test_browser(): assert not tctx.master.has_log("already running") b.browser.poll = lambda: None b.start() - assert tctx.master.has_log("already running") + assert await tctx.master.await_log("already running") b.done() assert not b.browser -def test_no_browser(): +@pytest.mark.asyncio +async def test_no_browser(): with mock.patch("shutil.which") as which: which.return_value = False b = browser.Browser() with taddons.context() as tctx: b.start() - assert tctx.master.has_log("platform is not supported") + assert await tctx.master.await_log("platform is not supported") diff --git a/test/mitmproxy/addons/test_cut.py b/test/mitmproxy/addons/test_cut.py index 56568f21..35375393 100644 --- a/test/mitmproxy/addons/test_cut.py +++ b/test/mitmproxy/addons/test_cut.py @@ -71,7 +71,8 @@ def qr(f): return fp.read() -def test_cut_clip(): +@pytest.mark.asyncio +async def test_cut_clip(): v = view.View() c = cut.Cut() with taddons.context() as tctx: @@ -95,7 +96,7 @@ def test_cut_clip(): "copy/paste mechanism for your system." pc.side_effect = pyperclip.PyperclipException(log_message) tctx.command(c.clip, "@all", "request.method") - assert tctx.master.has_log(log_message, level="error") + assert await tctx.master.await_log(log_message, level="error") def test_cut_save(tmpdir): @@ -125,7 +126,8 @@ def test_cut_save(tmpdir): (IsADirectoryError, "Is a directory"), (FileNotFoundError, "No such file or directory") ]) -def test_cut_save_open(exception, log_message, tmpdir): +@pytest.mark.asyncio +async def test_cut_save_open(exception, log_message, tmpdir): f = str(tmpdir.join("path")) v = view.View() c = cut.Cut() @@ -136,7 +138,7 @@ def test_cut_save_open(exception, log_message, tmpdir): with mock.patch("mitmproxy.addons.cut.open") as m: m.side_effect = exception(log_message) tctx.command(c.save, "@all", "request.method", f) - assert tctx.master.has_log(log_message, level="error") + assert await tctx.master.await_log(log_message, level="error") def test_cut(): diff --git a/test/mitmproxy/addons/test_dumper.py b/test/mitmproxy/addons/test_dumper.py index 228bacf8..c24801e4 100644 --- a/test/mitmproxy/addons/test_dumper.py +++ b/test/mitmproxy/addons/test_dumper.py @@ -141,15 +141,16 @@ def test_echo_request_line(): class TestContentView: - @mock.patch("mitmproxy.contentviews.auto.ViewAuto.__call__") - def test_contentview(self, view_auto): - view_auto.side_effect = exceptions.ContentViewException("") - sio = io.StringIO() - d = dumper.Dumper(sio) - with taddons.context(d) as ctx: - ctx.configure(d, flow_detail=4) - d.response(tflow.tflow()) - assert ctx.master.has_log("content viewer failed") + @pytest.mark.asyncio + async def test_contentview(self): + with mock.patch("mitmproxy.contentviews.auto.ViewAuto.__call__") as va: + va.side_effect = exceptions.ContentViewException("") + sio = io.StringIO() + d = dumper.Dumper(sio) + with taddons.context(d) as ctx: + ctx.configure(d, flow_detail=4) + d.response(tflow.tflow()) + assert await ctx.master.await_log("content viewer failed") def test_tcp(): diff --git a/test/mitmproxy/addons/test_export.py b/test/mitmproxy/addons/test_export.py index b625df56..f4bb0f64 100644 --- a/test/mitmproxy/addons/test_export.py +++ b/test/mitmproxy/addons/test_export.py @@ -125,17 +125,19 @@ def test_export(tmpdir): (IsADirectoryError, "Is a directory"), (FileNotFoundError, "No such file or directory") ]) -def test_export_open(exception, log_message, tmpdir): +@pytest.mark.asyncio +async def test_export_open(exception, log_message, tmpdir): f = str(tmpdir.join("path")) e = export.Export() with taddons.context() as tctx: with mock.patch("mitmproxy.addons.export.open") as m: m.side_effect = exception(log_message) e.file("raw", tflow.tflow(resp=True), f) - assert tctx.master.has_log(log_message, level="error") + assert await tctx.master.await_log(log_message, level="error") -def test_clip(tmpdir): +@pytest.mark.asyncio +async def test_clip(tmpdir): e = export.Export() with taddons.context() as tctx: with pytest.raises(exceptions.CommandError): @@ -158,4 +160,4 @@ def test_clip(tmpdir): "copy/paste mechanism for your system." pc.side_effect = pyperclip.PyperclipException(log_message) e.clip("raw", tflow.tflow(resp=True)) - assert tctx.master.has_log(log_message, level="error") + assert await tctx.master.await_log(log_message, level="error") diff --git a/test/mitmproxy/addons/test_readfile.py b/test/mitmproxy/addons/test_readfile.py index 0439862a..a0826a26 100644 --- a/test/mitmproxy/addons/test_readfile.py +++ b/test/mitmproxy/addons/test_readfile.py @@ -55,26 +55,28 @@ class TestReadFile: with pytest.raises(exceptions.OptionsError): rf.running() - @mock.patch('mitmproxy.master.Master.load_flow') - def test_corrupt(self, mck, corrupt_data): + @pytest.mark.asyncio + async def test_corrupt(self, corrupt_data): rf = readfile.ReadFile() with taddons.context(rf) as tctx: - with pytest.raises(exceptions.FlowReadException): - rf.load_flows(io.BytesIO(b"qibble")) - assert not mck.called - assert len(tctx.master.logs) == 1 - - with pytest.raises(exceptions.FlowReadException): - rf.load_flows(corrupt_data) - assert mck.called - assert len(tctx.master.logs) == 2 - - def test_nonexisting_file(self): + with mock.patch('mitmproxy.master.Master.load_flow') as mck: + with pytest.raises(exceptions.FlowReadException): + rf.load_flows(io.BytesIO(b"qibble")) + assert not mck.called + + tctx.master.clear() + with pytest.raises(exceptions.FlowReadException): + rf.load_flows(corrupt_data) + assert await tctx.master.await_log("file corrupted") + assert mck.called + + @pytest.mark.asyncio + async def test_nonexisting_file(self): rf = readfile.ReadFile() with taddons.context(rf) as tctx: with pytest.raises(exceptions.FlowReadException): rf.load_flows_from_path("nonexistent") - assert len(tctx.master.logs) == 1 + assert await tctx.master.await_log("nonexistent") class TestReadFileStdin: diff --git a/test/mitmproxy/addons/test_replace.py b/test/mitmproxy/addons/test_replace.py index 9c1f7f79..dce58ae0 100644 --- a/test/mitmproxy/addons/test_replace.py +++ b/test/mitmproxy/addons/test_replace.py @@ -79,7 +79,8 @@ class TestReplaceFile: r.request(f) assert f.request.content == b"bar" - def test_nonexistent(self, tmpdir): + @pytest.mark.asyncio + async def test_nonexistent(self, tmpdir): r = replace.Replace() with taddons.context(r) as tctx: with pytest.raises(Exception, match="Invalid file path"): @@ -97,6 +98,5 @@ class TestReplaceFile: tmpfile.remove() f = tflow.tflow() f.request.content = b"foo" - assert not tctx.master.logs r.request(f) - assert tctx.master.logs + assert await tctx.master.await_log("could not read") diff --git a/test/mitmproxy/addons/test_script.py b/test/mitmproxy/addons/test_script.py index 5427c788..17b1e372 100644 --- a/test/mitmproxy/addons/test_script.py +++ b/test/mitmproxy/addons/test_script.py @@ -1,14 +1,11 @@ import os import sys import traceback -from unittest import mock -import time import pytest from mitmproxy import addonmanager from mitmproxy import exceptions -from mitmproxy import log from mitmproxy.addons import script from mitmproxy.test import taddons from mitmproxy.test import tflow @@ -58,7 +55,7 @@ async def test_script_print_stdout(): tutils.test_data.path("mitmproxy/data/addonscripts/print.py") ) ns.load(addonmanager.Loader(tctx.master)) - assert tctx.master.has_log("stdoutprint") + assert await tctx.master.await_log("stdoutprint") class TestScript: @@ -100,7 +97,8 @@ class TestScript: assert rec.call_log[0][1] == "request" - def test_reload(self, tmpdir): + @pytest.mark.asyncio + async def test_reload(self, tmpdir): with taddons.context() as tctx: f = tmpdir.join("foo.py") f.ensure(file=True) @@ -108,15 +106,15 @@ class TestScript: sc = script.Script(str(f)) tctx.configure(sc) sc.tick() - assert tctx.master.has_log("Loading") + assert await tctx.master.await_log("Loading") tctx.master.clear() - assert not tctx.master.has_log("Loading") sc.last_load, sc.last_mtime = 0, 0 sc.tick() - assert tctx.master.has_log("Loading") + assert await tctx.master.await_log("Loading") - def test_exception(self): + @pytest.mark.asyncio + async def test_exception(self): with taddons.context() as tctx: sc = script.Script( tutils.test_data.path("mitmproxy/data/addonscripts/error.py") @@ -128,8 +126,8 @@ class TestScript: f = tflow.tflow(resp=True) tctx.master.addons.trigger("request", f) - assert tctx.master.has_log("ValueError: Error!") - assert tctx.master.has_log("error.py") + assert await tctx.master.await_log("ValueError: Error!") + assert await tctx.master.await_log("error.py") def test_addon(self): with taddons.context() as tctx: @@ -165,13 +163,15 @@ class TestCutTraceback: class TestScriptLoader: - def test_script_run(self): + @pytest.mark.asyncio + async def test_script_run(self): rp = tutils.test_data.path( "mitmproxy/data/addonscripts/recorder/recorder.py" ) sc = script.ScriptLoader() with taddons.context(sc) as tctx: sc.script_run([tflow.tflow(resp=True)], rp) + await tctx.master.await_log("recorder response") debug = [i.msg for i in tctx.master.logs if i.level == "debug"] assert debug == [ 'recorder load', 'recorder running', 'recorder configure', @@ -242,19 +242,21 @@ class TestScriptLoader: tctx.invoke(sc, "tick") assert len(tctx.master.addons) == 1 - def test_script_error_handler(self): + @pytest.mark.asyncio + async def test_script_error_handler(self): path = "/sample/path/example.py" exc = SyntaxError msg = "Error raised" tb = True with taddons.context() as tctx: script.script_error_handler(path, exc, msg, tb) - assert tctx.master.has_log("/sample/path/example.py") - assert tctx.master.has_log("Error raised") - assert tctx.master.has_log("lineno") - assert tctx.master.has_log("NoneType") + assert await tctx.master.await_log("/sample/path/example.py") + assert await tctx.master.await_log("Error raised") + assert await tctx.master.await_log("lineno") + assert await tctx.master.await_log("NoneType") - def test_order(self): + @pytest.mark.asyncio + async def test_order(self): rec = tutils.test_data.path("mitmproxy/data/addonscripts/recorder") sc = script.ScriptLoader() sc.is_running = True @@ -268,6 +270,7 @@ class TestScriptLoader: ] ) tctx.master.addons.invoke_addon(sc, "tick") + await tctx.master.await_log("c tick") debug = [i.msg for i in tctx.master.logs if i.level == "debug"] assert debug == [ 'a load', @@ -286,7 +289,7 @@ class TestScriptLoader: 'c tick', ] - tctx.master.logs = [] + tctx.master.clear() tctx.configure( sc, scripts = [ @@ -296,6 +299,7 @@ class TestScriptLoader: ] ) + await tctx.master.await_log("c configure") debug = [i.msg for i in tctx.master.logs if i.level == "debug"] assert debug == [ 'c configure', @@ -312,6 +316,7 @@ class TestScriptLoader: ] ) tctx.master.addons.invoke_addon(sc, "tick") + await tctx.master.await_log("a tick") debug = [i.msg for i in tctx.master.logs if i.level == "debug"] assert debug == [ diff --git a/test/mitmproxy/addons/test_termstatus.py b/test/mitmproxy/addons/test_termstatus.py index 5f960a1c..c4802102 100644 --- a/test/mitmproxy/addons/test_termstatus.py +++ b/test/mitmproxy/addons/test_termstatus.py @@ -1,15 +1,17 @@ +import pytest + from mitmproxy import proxy from mitmproxy.addons import termstatus from mitmproxy.test import taddons -def test_configure(): +@pytest.mark.asyncio +async def test_configure(): ts = termstatus.TermStatus() with taddons.context() as ctx: ctx.master.server = proxy.DummyServer() ctx.configure(ts, server=False) ts.running() - assert not ctx.master.logs ctx.configure(ts, server=True) ts.running() - assert ctx.master.logs + await ctx.master.await_log("server listening") diff --git a/test/mitmproxy/addons/test_view.py b/test/mitmproxy/addons/test_view.py index 6f2a9ca5..62a6aeb0 100644 --- a/test/mitmproxy/addons/test_view.py +++ b/test/mitmproxy/addons/test_view.py @@ -159,7 +159,8 @@ def test_orders(): assert v.order_options() -def test_load(tmpdir): +@pytest.mark.asyncio +async def test_load(tmpdir): path = str(tmpdir.join("path")) v = view.View() with taddons.context() as tctx: @@ -182,7 +183,7 @@ def test_load(tmpdir): with open(path, "wb") as f: f.write(b"invalidflows") v.load_file(path) - assert tctx.master.has_log("Invalid data format.") + assert await tctx.master.await_log("Invalid data format.") def test_resolve(): diff --git a/test/mitmproxy/proxy/protocol/test_websocket.py b/test/mitmproxy/proxy/protocol/test_websocket.py index e5ed8e9d..352a9c0d 100644 --- a/test/mitmproxy/proxy/protocol/test_websocket.py +++ b/test/mitmproxy/proxy/protocol/test_websocket.py @@ -3,7 +3,6 @@ import os import struct import tempfile import traceback -import time from mitmproxy import options from mitmproxy import exceptions @@ -314,7 +313,8 @@ class TestPong(_WebSocketTest): wfile.flush() websockets.Frame.from_file(rfile) - def test_pong(self): + @pytest.mark.asyncio + async def test_pong(self): self.setup_connection() self.client.wfile.write(bytes(websockets.Frame(fin=1, mask=1, opcode=websockets.OPCODE.PING, payload=b'foobar'))) @@ -327,12 +327,7 @@ class TestPong(_WebSocketTest): assert frame.header.opcode == websockets.OPCODE.PONG assert frame.payload == b'foobar' - for i in range(20): - if self.master.has_log("Pong Received from server", "info"): - break - time.sleep(0.01) - else: - raise AssertionError("No pong seen") + assert await self.master.await_log("pong received") class TestClose(_WebSocketTest): diff --git a/test/mitmproxy/script/test_concurrent.py b/test/mitmproxy/script/test_concurrent.py index ceff9fb9..876093f7 100644 --- a/test/mitmproxy/script/test_concurrent.py +++ b/test/mitmproxy/script/test_concurrent.py @@ -1,3 +1,5 @@ +import pytest + from mitmproxy.test import tflow from mitmproxy.test import tutils from mitmproxy.test import taddons @@ -31,14 +33,15 @@ class TestConcurrent(tservers.MasterTest): return raise ValueError("Script never acked") - def test_concurrent_err(self): + @pytest.mark.asyncio + async def test_concurrent_err(self): with taddons.context() as tctx: tctx.script( tutils.test_data.path( "mitmproxy/data/addonscripts/concurrent_decorator_err.py" ) ) - assert tctx.master.has_log("decorator not supported") + assert await tctx.master.await_log("decorator not supported") def test_concurrent_class(self): with taddons.context() as tctx: diff --git a/test/mitmproxy/test_addonmanager.py b/test/mitmproxy/test_addonmanager.py index 4915cc36..e8e3a5f7 100644 --- a/test/mitmproxy/test_addonmanager.py +++ b/test/mitmproxy/test_addonmanager.py @@ -87,7 +87,8 @@ def test_defaults(): assert addons.default_addons() -def test_loader(): +@pytest.mark.asyncio +async def test_loader(): with taddons.context() as tctx: l = addonmanager.Loader(tctx.master) l.add_option("custom_option", bool, False, "help") @@ -99,7 +100,7 @@ def test_loader(): # a different signature should emit a warning though. l.add_option("custom_option", bool, True, "help") - assert tctx.master.has_log("Over-riding existing option") + assert await tctx.master.await_log("Over-riding existing option") def cmd(a: str) -> str: return "foo" @@ -107,7 +108,8 @@ def test_loader(): l.add_command("test.command", cmd) -def test_simple(): +@pytest.mark.asyncio +async def test_simple(): with taddons.context(loadcore=False) as tctx: a = tctx.master.addons @@ -121,14 +123,14 @@ def test_simple(): assert not a.chain a.add(TAddon("one")) - a.trigger("done") + a.trigger("running") a.trigger("tick") - assert tctx.master.has_log("not callable") + assert await tctx.master.await_log("not callable") tctx.master.clear() a.get("one").tick = addons a.trigger("tick") - assert not tctx.master.has_log("not callable") + assert not await tctx.master.await_log("not callable") a.remove(a.get("one")) assert not a.get("one") diff --git a/test/mitmproxy/test_controller.py b/test/mitmproxy/test_controller.py index f7c64ed9..953d37e1 100644 --- a/test/mitmproxy/test_controller.py +++ b/test/mitmproxy/test_controller.py @@ -5,12 +5,11 @@ import pytest from mitmproxy.exceptions import Kill, ControlException from mitmproxy import controller from mitmproxy.test import taddons +import mitmproxy.ctx @pytest.mark.asyncio async def test_master(): - class TMsg: - pass class tAddon: def log(self, _): @@ -20,12 +19,11 @@ async def test_master(): assert not ctx.master.should_exit.is_set() async def test(): - msg = TMsg() - msg.reply = controller.DummyReply() - await ctx.master.channel.tell("log", msg) + mitmproxy.ctx.log("test") asyncio.ensure_future(test()) - assert not ctx.master.should_exit.is_set() + assert await ctx.master.await_log("test") + assert ctx.master.should_exit.is_set() class TestReply: diff --git a/test/mitmproxy/test_taddons.py b/test/mitmproxy/test_taddons.py index 5a4c99fc..67f64463 100644 --- a/test/mitmproxy/test_taddons.py +++ b/test/mitmproxy/test_taddons.py @@ -1,21 +1,27 @@ import io + +import pytest + from mitmproxy.test import taddons from mitmproxy.test import tutils from mitmproxy import ctx -def test_recordingmaster(): +@pytest.mark.asyncio +async def test_recordingmaster(): with taddons.context() as tctx: assert not tctx.master.has_log("nonexistent") assert not tctx.master.has_event("nonexistent") ctx.log.error("foo") assert not tctx.master.has_log("foo", level="debug") - assert tctx.master.has_log("foo", level="error") + assert await tctx.master.await_log("foo", level="error") -def test_dumplog(): +@pytest.mark.asyncio +async def test_dumplog(): with taddons.context() as tctx: ctx.log.info("testing") + await ctx.master.await_log("testing") s = io.StringIO() tctx.master.dump_log(s) assert s.getvalue() diff --git a/test/mitmproxy/tools/console/test_defaultkeys.py b/test/mitmproxy/tools/console/test_defaultkeys.py index b09f245d..52075c84 100644 --- a/test/mitmproxy/tools/console/test_defaultkeys.py +++ b/test/mitmproxy/tools/console/test_defaultkeys.py @@ -6,6 +6,7 @@ from mitmproxy import command import pytest + @pytest.mark.asyncio async def test_commands_exist(): km = keymap.Keymap(None) diff --git a/test/mitmproxy/tools/console/test_master.py b/test/mitmproxy/tools/console/test_master.py index a4318e22..ee35e87f 100644 --- a/test/mitmproxy/tools/console/test_master.py +++ b/test/mitmproxy/tools/console/test_master.py @@ -6,9 +6,8 @@ from ... import tservers import pytest -@pytest.mark.asyncio - +@pytest.mark.asyncio class TestMaster(tservers.MasterTest): def mkmaster(self, **opts): o = options.Options(**opts) -- cgit v1.2.3 From ea213782d8adaa21aa99f2de818172ee872cf3e9 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Wed, 4 Apr 2018 14:25:37 +1200 Subject: asyncio: remove test master.has_log Now that logs are async, using this call is almost always a mistake. Signal this by making it semi-private. The method may go away entirely down the track. --- mitmproxy/test/taddons.py | 4 ++-- test/mitmproxy/addons/test_browser.py | 3 +-- test/mitmproxy/addons/test_check_ca.py | 9 ++++++--- test/mitmproxy/proxy/protocol/test_websocket.py | 5 +++-- test/mitmproxy/proxy/test_server.py | 18 ++++++++++------- test/mitmproxy/test_addonmanager.py | 27 ++++++++++++++----------- test/mitmproxy/test_taddons.py | 4 ++-- 7 files changed, 40 insertions(+), 30 deletions(-) diff --git a/mitmproxy/test/taddons.py b/mitmproxy/test/taddons.py index 5cef80f1..6d439c4a 100644 --- a/mitmproxy/test/taddons.py +++ b/mitmproxy/test/taddons.py @@ -35,7 +35,7 @@ class RecordingMaster(mitmproxy.master.Master): for i in self.logs: print("%s: %s" % (i.level, i.msg), file=outf) - def has_log(self, txt, level=None): + def _has_log(self, txt, level=None): for i in self.logs: if level and i.level != level: continue @@ -45,7 +45,7 @@ class RecordingMaster(mitmproxy.master.Master): async def await_log(self, txt, level=None): for i in range(20): - if self.has_log(txt, level): + if self._has_log(txt, level): return True else: await asyncio.sleep(0.1) diff --git a/test/mitmproxy/addons/test_browser.py b/test/mitmproxy/addons/test_browser.py index 041bab04..b05996fe 100644 --- a/test/mitmproxy/addons/test_browser.py +++ b/test/mitmproxy/addons/test_browser.py @@ -13,9 +13,8 @@ async def test_browser(): with taddons.context() as tctx: b.start() assert po.called - b.start() - assert not tctx.master.has_log("already running") + b.start() b.browser.poll = lambda: None b.start() assert await tctx.master.await_log("already running") diff --git a/test/mitmproxy/addons/test_check_ca.py b/test/mitmproxy/addons/test_check_ca.py index cd34a9be..5e820b6d 100644 --- a/test/mitmproxy/addons/test_check_ca.py +++ b/test/mitmproxy/addons/test_check_ca.py @@ -8,12 +8,15 @@ from mitmproxy.test import taddons class TestCheckCA: @pytest.mark.parametrize('expired', [False, True]) - def test_check_ca(self, expired): + @pytest.mark.asyncio + async def test_check_ca(self, expired): msg = 'The mitmproxy certificate authority has expired!' with taddons.context() as tctx: tctx.master.server = mock.MagicMock() - tctx.master.server.config.certstore.default_ca.has_expired = mock.MagicMock(return_value=expired) + tctx.master.server.config.certstore.default_ca.has_expired = mock.MagicMock( + return_value = expired + ) a = check_ca.CheckCA() tctx.configure(a) - assert tctx.master.has_log(msg) is expired + assert await tctx.master.await_log(msg) == expired diff --git a/test/mitmproxy/proxy/protocol/test_websocket.py b/test/mitmproxy/proxy/protocol/test_websocket.py index 352a9c0d..347c0dfc 100644 --- a/test/mitmproxy/proxy/protocol/test_websocket.py +++ b/test/mitmproxy/proxy/protocol/test_websocket.py @@ -285,7 +285,8 @@ class TestPing(_WebSocketTest): wfile.flush() websockets.Frame.from_file(rfile) - def test_ping(self): + @pytest.mark.asyncio + async def test_ping(self): self.setup_connection() frame = websockets.Frame.from_file(self.client.rfile) @@ -295,7 +296,7 @@ class TestPing(_WebSocketTest): assert frame.header.opcode == websockets.OPCODE.PING assert frame.payload == b'' # We don't send payload to other end - assert self.master.has_log("Pong Received from server", "info") + assert await self.master.await_log("Pong Received from server", "info") class TestPong(_WebSocketTest): diff --git a/test/mitmproxy/proxy/test_server.py b/test/mitmproxy/proxy/test_server.py index 31033e25..bf24e28b 100644 --- a/test/mitmproxy/proxy/test_server.py +++ b/test/mitmproxy/proxy/test_server.py @@ -1,3 +1,4 @@ +import asyncio import os import socket import time @@ -234,13 +235,14 @@ class TestHTTP(tservers.HTTPProxyTest, CommonMixin): assert p.request(req) assert p.request(req) - def test_get_connection_switching(self): + @pytest.mark.asyncio + async def test_get_connection_switching(self): req = "get:'%s/p/200:b@1'" p = self.pathoc() with p.connect(): assert p.request(req % self.server.urlbase) assert p.request(req % self.server2.urlbase) - assert self.proxy.tmaster.has_log("serverdisconnect") + assert await self.proxy.tmaster.await_log("serverdisconnect") def test_blank_leading_line(self): p = self.pathoc() @@ -443,13 +445,14 @@ class TestReverse(tservers.ReverseProxyTest, CommonMixin, TcpMixin): req = self.master.state.flows[0].request assert req.host_header == "127.0.0.1" - def test_selfconnection(self): + @pytest.mark.asyncio + async def test_selfconnection(self): self.options.mode = "reverse:http://127.0.0.1:0" p = self.pathoc() with p.connect(): p.request("get:/") - assert self.master.has_log("The proxy shall not connect to itself.") + assert await self.master.await_log("The proxy shall not connect to itself.") class TestReverseSSL(tservers.ReverseProxyTest, CommonMixin, TcpMixin): @@ -549,7 +552,6 @@ class TestHttps2Http(tservers.ReverseProxyTest): p = self.pathoc(ssl=True, sni="example.com") with p.connect(): assert p.request("get:'/p/200'").status_code == 200 - assert not self.proxy.tmaster.has_log("error in handle_sni") def test_http(self): p = self.pathoc(ssl=False) @@ -814,11 +816,13 @@ class TestServerConnect(tservers.HTTPProxyTest): opts.upstream_cert = False return opts - def test_unnecessary_serverconnect(self): + @pytest.mark.asyncio + async def test_unnecessary_serverconnect(self): """A replayed/fake response with no upstream_cert should not connect to an upstream server""" self.set_addons(AFakeResponse()) assert self.pathod("200").status_code == 200 - assert not self.proxy.tmaster.has_log("serverconnect") + asyncio.sleep(0.1) + assert not self.proxy.tmaster._has_log("serverconnect") class AKillRequest: diff --git a/test/mitmproxy/test_addonmanager.py b/test/mitmproxy/test_addonmanager.py index e8e3a5f7..796ae1bd 100644 --- a/test/mitmproxy/test_addonmanager.py +++ b/test/mitmproxy/test_addonmanager.py @@ -1,4 +1,6 @@ import pytest +from unittest import mock + from mitmproxy import addons from mitmproxy import addonmanager @@ -90,22 +92,23 @@ def test_defaults(): @pytest.mark.asyncio async def test_loader(): with taddons.context() as tctx: - l = addonmanager.Loader(tctx.master) - l.add_option("custom_option", bool, False, "help") - assert "custom_option" in l.master.options + with mock.patch("mitmproxy.ctx.log.warn") as warn: + l = addonmanager.Loader(tctx.master) + l.add_option("custom_option", bool, False, "help") + assert "custom_option" in l.master.options - # calling this again with the same signature is a no-op. - l.add_option("custom_option", bool, False, "help") - assert not tctx.master.has_log("Over-riding existing option") + # calling this again with the same signature is a no-op. + l.add_option("custom_option", bool, False, "help") + assert not warn.called - # a different signature should emit a warning though. - l.add_option("custom_option", bool, True, "help") - assert await tctx.master.await_log("Over-riding existing option") + # a different signature should emit a warning though. + l.add_option("custom_option", bool, True, "help") + assert warn.called - def cmd(a: str) -> str: - return "foo" + def cmd(a: str) -> str: + return "foo" - l.add_command("test.command", cmd) + l.add_command("test.command", cmd) @pytest.mark.asyncio diff --git a/test/mitmproxy/test_taddons.py b/test/mitmproxy/test_taddons.py index 67f64463..39dd65c2 100644 --- a/test/mitmproxy/test_taddons.py +++ b/test/mitmproxy/test_taddons.py @@ -10,10 +10,10 @@ from mitmproxy import ctx @pytest.mark.asyncio async def test_recordingmaster(): with taddons.context() as tctx: - assert not tctx.master.has_log("nonexistent") + assert not tctx.master._has_log("nonexistent") assert not tctx.master.has_event("nonexistent") ctx.log.error("foo") - assert not tctx.master.has_log("foo", level="debug") + assert not tctx.master._has_log("foo", level="debug") assert await tctx.master.await_log("foo", level="error") -- cgit v1.2.3 From 6a08ef465f3af6cb2044baf98279de1071538f5a Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Wed, 4 Apr 2018 14:46:29 +1200 Subject: asyncio: remove master.add_log, in favor of a persistent log.Log instance --- mitmproxy/addonmanager.py | 12 +++--------- mitmproxy/addons/check_ca.py | 4 ++-- mitmproxy/master.py | 9 ++------- mitmproxy/tools/console/flowview.py | 2 +- mitmproxy/tools/web/master.py | 8 +++----- test/mitmproxy/tools/web/test_app.py | 2 +- 6 files changed, 12 insertions(+), 25 deletions(-) diff --git a/mitmproxy/addonmanager.py b/mitmproxy/addonmanager.py index 418ab046..7c446c80 100644 --- a/mitmproxy/addonmanager.py +++ b/mitmproxy/addonmanager.py @@ -41,8 +41,8 @@ class StreamLog: """ A class for redirecting output using contextlib. """ - def __init__(self, log): - self.log = log + def __init__(self, lg): + self.log = lg def write(self, buf): if buf.strip(): @@ -55,13 +55,7 @@ class StreamLog: @contextlib.contextmanager def safecall(): - # resolve ctx.master here. - # we want to be threadsafe, and ctx.master may already be cleared when an addon prints(). - channel = ctx.master.channel - # don't use master.add_log (which is not thread-safe). Instead, put on event queue. - stdout_replacement = StreamLog( - lambda message: channel("log", log.LogEntry(message, "warn")) - ) + stdout_replacement = StreamLog(lambda message: ctx.log.warn(message)) try: with contextlib.redirect_stdout(stdout_replacement): yield diff --git a/mitmproxy/addons/check_ca.py b/mitmproxy/addons/check_ca.py index f786af5a..447ba64d 100644 --- a/mitmproxy/addons/check_ca.py +++ b/mitmproxy/addons/check_ca.py @@ -1,4 +1,5 @@ import mitmproxy +from mitmproxy import ctx class CheckCA: @@ -15,10 +16,9 @@ class CheckCA: if has_ca: self.failed = mitmproxy.ctx.master.server.config.certstore.default_ca.has_expired() if self.failed: - mitmproxy.ctx.master.add_log( + ctx.log.warn( "The mitmproxy certificate authority has expired!\n" "Please delete all CA-related files in your ~/.mitmproxy folder.\n" "The CA will be regenerated automatically after restarting mitmproxy.\n" "Then make sure all your clients have the new CA installed.", - "warn", ) diff --git a/mitmproxy/master.py b/mitmproxy/master.py index be639dcd..323f73dc 100644 --- a/mitmproxy/master.py +++ b/mitmproxy/master.py @@ -56,6 +56,7 @@ class Master: self._server = None self.first_tick = True self.waiting_flows = [] + self.log = log.Log(self) @property def server(self): @@ -73,7 +74,7 @@ class Master: yield return mitmproxy_ctx.master = self - mitmproxy_ctx.log = log.Log(self) + mitmproxy_ctx.log = self.log mitmproxy_ctx.options = self.options try: yield @@ -82,12 +83,6 @@ class Master: mitmproxy_ctx.log = None mitmproxy_ctx.options = None - def add_log(self, e, level): - """ - level: debug, alert, info, warn, error - """ - self.addons.trigger("log", log.LogEntry(e, level)) - def start(self): self.should_exit.clear() if self.server: diff --git a/mitmproxy/tools/console/flowview.py b/mitmproxy/tools/console/flowview.py index 650a9366..837d71a8 100644 --- a/mitmproxy/tools/console/flowview.py +++ b/mitmproxy/tools/console/flowview.py @@ -121,7 +121,7 @@ class FlowDetails(tabs.Tabs): viewmode, message ) if error: - self.master.add_log(error, "debug") + self.master.log.debug(error) # Give hint that you have to tab for the response. if description == "No content" and isinstance(message, http.HTTPRequest): description = "No request content (press tab to view response)" diff --git a/mitmproxy/tools/web/master.py b/mitmproxy/tools/web/master.py index b7eddcce..575694c7 100644 --- a/mitmproxy/tools/web/master.py +++ b/mitmproxy/tools/web/master.py @@ -114,17 +114,15 @@ class WebMaster(master.Master): iol.add_callback(self.start) web_url = "http://{}:{}/".format(self.options.web_iface, self.options.web_port) - self.add_log( - "Web server listening at {}".format(web_url), - "info" + self.log.info( + "Web server listening at {}".format(web_url), ) if self.options.web_open_browser: success = open_browser(web_url) if not success: - self.add_log( + self.log.info( "No web browser found. Please open a browser and point it to {}".format(web_url), - "info" ) try: iol.start() diff --git a/test/mitmproxy/tools/web/test_app.py b/test/mitmproxy/tools/web/test_app.py index 994690ea..00188895 100644 --- a/test/mitmproxy/tools/web/test_app.py +++ b/test/mitmproxy/tools/web/test_app.py @@ -45,7 +45,7 @@ class TestApp(tornado.testing.AsyncHTTPTestCase): f.id = "42" m.view.add([f]) m.view.add([tflow.tflow(err=True)]) - m.add_log("test log", "info") + m.log.info("test log") self.master = m self.view = m.view self.events = m.events -- cgit v1.2.3 From 9dcd15d350aa5e2ee2d6f7299d0183cabddeb992 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Wed, 4 Apr 2018 14:52:24 +1200 Subject: asyncio: ditch the handler context There are a few reasons for this. First, logs are now async, and can be called at any time. Second, the event loop is thread local, so there can only ever be one master per thread. These two things together completely obviate the need for a handler context. --- mitmproxy/addonmanager.py | 30 +++++++++---------------- mitmproxy/command.py | 6 +---- mitmproxy/master.py | 21 ++++------------- mitmproxy/test/taddons.py | 11 --------- setup.py | 1 + test/mitmproxy/addons/test_script.py | 5 +++-- 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 diff --git a/setup.py b/setup.py index 18fb3827..3e914847 100644 --- a/setup.py +++ b/setup.py @@ -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): -- cgit v1.2.3