Grasscutters/mitmproxy
3
0
Fork 0
mirror of https://github.com/Grasscutters/mitmproxy.git synced 2024-11-23 00:01:36 +00:00
Code Issues Projects Releases Wiki Activity

asyncio: move log mechanism onto the event loop

Browse Source 
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.
This commit is contained in:
Aldo Cortesi 2018-04-04 07:28:43 +12:00 committed by Aldo Cortesi
parent 0fa1280daa
commit 80f2bac356
21 changed files with 126 additions and 92 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
mitmproxy/addonmanager.py
View File

@ -3,7 +3,6 @@ import typing
import traceback
import contextlib
import sys
import asyncio
from mitmproxy import exceptions
from mitmproxy import eventsequence

6
mitmproxy/log.py
View File

@ -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 = [

1
mitmproxy/master.py
View File

@ -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)

9
mitmproxy/test/taddons.py
View File

@ -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:

5
test/mitmproxy/addons/test_allowremote.py
View File

@ -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()

11
test/mitmproxy/addons/test_browser.py
View File

@ -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")

10
test/mitmproxy/addons/test_cut.py
View File

@ -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():

19
test/mitmproxy/addons/test_dumper.py
View File

@ -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():

10
test/mitmproxy/addons/test_export.py
View File

@ -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")

26
test/mitmproxy/addons/test_readfile.py
View File

@ -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 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
with pytest.raises(exceptions.FlowReadException):
rf.load_flows(corrupt_data)
assert mck.called
assert len(tctx.master.logs) == 2
tctx.master.clear()
with pytest.raises(exceptions.FlowReadException):
rf.load_flows(corrupt_data)
assert await tctx.master.await_log("file corrupted")
assert mck.called
def test_nonexisting_file(self):
@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:

6
test/mitmproxy/addons/test_replace.py
View File

@ -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")

43
test/mitmproxy/addons/test_script.py
View File

@ -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 == [

8
test/mitmproxy/addons/test_termstatus.py
View File

@ -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")

5
test/mitmproxy/addons/test_view.py
View File

@ -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():

11
test/mitmproxy/proxy/protocol/test_websocket.py
View File

@ -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):

7
test/mitmproxy/script/test_concurrent.py
View File

@ -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:

14
test/mitmproxy/test_addonmanager.py
View File

@ -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")

10
test/mitmproxy/test_controller.py
View File

@ -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:

12
test/mitmproxy/test_taddons.py
View File

@ -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()

1
test/mitmproxy/tools/console/test_defaultkeys.py
View File

@ -6,6 +6,7 @@ from mitmproxy import command
import pytest
@pytest.mark.asyncio
async def test_commands_exist():
km = keymap.Keymap(None)

3
test/mitmproxy/tools/console/test_master.py
View File

@ -6,9 +6,8 @@ from ... import tservers
import pytest
@pytest.mark.asyncio
class TestMaster(tservers.MasterTest):
def mkmaster(self, **opts):
o = options.Options(**opts)