Remove test handlers by using taddons.RecordingMaster

This also means expanding and tweaking the recording master API, which we
reflect through the current test suite
This commit is contained in:
Aldo Cortesi 2017-03-16 12:50:33 +13:00 committed by Aldo Cortesi
parent 85ddc5056b
commit 1410cbb4b6
13 changed files with 64 additions and 119 deletions

View File

@ -1,43 +0,0 @@
#!/usr/bin/env python3
"""
This example shows how to build a proxy based on mitmproxy's Flow
primitives.
Heads Up: In the majority of cases, you want to use inline scripts.
Note that request and response messages are not automatically replied to,
so we need to implement handlers to do this.
"""
from mitmproxy import controller, options, master
from mitmproxy.proxy import ProxyServer, ProxyConfig
class MyMaster(master.Master):
def run(self):
try:
master.Master.run(self)
except KeyboardInterrupt:
self.shutdown()
@controller.handler
def request(self, f):
print("request", f)
@controller.handler
def response(self, f):
print("response", f)
@controller.handler
def error(self, f):
print("error", f)
@controller.handler
def log(self, l):
print("log", l.msg)
opts = options.Options(cadir="~/.mitmproxy/")
config = ProxyConfig(opts)
server = ProxyServer(config)
m = MyMaster(opts, server)
m.run()

View File

@ -1,5 +1,6 @@
import mitmproxy
from mitmproxy.net import tcp
from mitmproxy import ctx
class CheckALPN:
@ -9,9 +10,8 @@ class CheckALPN:
def configure(self, options, updated):
self.failed = mitmproxy.ctx.master.options.http2 and not tcp.HAS_ALPN
if self.failed:
mitmproxy.ctx.master.add_log(
ctx.log.warn(
"HTTP/2 is disabled because ALPN support missing!\n"
"OpenSSL 1.0.2+ required to support HTTP/2 connections.\n"
"Use --no-http2 to silence this warning.",
"warn",
"Use --no-http2 to silence this warning."
)

View File

@ -12,7 +12,10 @@ class _AddonWrapper:
self.addons = addons
def trigger(self, event, *args, **kwargs):
self.master.events.append((event, args, kwargs))
if event == "log":
self.master.logs.append(args[0])
else:
self.master.events.append((event, args, kwargs))
return self.addons.trigger(event, *args, **kwargs)
def __getattr__(self, attr):
@ -26,13 +29,19 @@ class RecordingMaster(mitmproxy.master.Master):
self.events = []
self.logs = []
def has_log(self, txt, level=None):
for i in self.logs:
if level and i.level != level:
continue
if txt.lower() in i.msg.lower():
return True
return False
def has_event(self, name):
for i in self.events:
if i[0] == name:
return True
def add_log(self, e, level):
self.logs.append((level, e))
return False
def clear(self):
self.logs = []

View File

@ -13,7 +13,6 @@ import traceback
import urwid
from mitmproxy import addons
from mitmproxy import controller
from mitmproxy import exceptions
from mitmproxy import master
from mitmproxy import io

View File

@ -1,4 +1,3 @@
from mitmproxy import ctx
from mitmproxy import addons
from mitmproxy import options
from mitmproxy import master

View File

@ -12,7 +12,7 @@ class TestCheckALPN:
with taddons.context() as tctx:
a = check_alpn.CheckALPN()
tctx.configure(a)
assert not any(msg in m for l, m in tctx.master.logs)
assert not tctx.master.has_log(msg)
def test_check_no_alpn(self, disable_alpn):
msg = 'ALPN support missing'
@ -20,4 +20,4 @@ class TestCheckALPN:
with taddons.context() as tctx:
a = check_alpn.CheckALPN()
tctx.configure(a)
assert any(msg in m for l, m in tctx.master.logs)
assert tctx.master.has_log(msg)

View File

@ -16,4 +16,4 @@ class TestCheckCA:
tctx.master.server.config.certstore.default_ca.has_expired = mock.MagicMock(return_value=expired)
a = check_ca.CheckCA()
tctx.configure(a)
assert any(msg in m for l, m in tctx.master.logs) is expired
assert tctx.master.has_log(msg) is expired

View File

@ -151,7 +151,7 @@ class TestContentView:
with taddons.context(options=options.Options()) as ctx:
ctx.configure(d, flow_detail=4, verbosity=3)
d.response(tflow.tflow())
assert "Content viewer failed" in ctx.master.logs[0][1]
assert ctx.master.has_log("content viewer failed")
def test_tcp():

View File

@ -22,14 +22,12 @@ def test_scriptenv():
with taddons.context() as tctx:
with script.scriptenv("path", []):
raise SystemExit
assert tctx.master.logs[0][0] == "error"
assert "exited" in tctx.master.logs[0][1]
assert tctx.master.has_log("exited", "error")
tctx.master.clear()
with script.scriptenv("path", []):
raise ValueError("fooo")
assert tctx.master.logs[0][0] == "error"
assert "foo" in tctx.master.logs[0][1]
assert tctx.master.has_log("fooo", "error")
class Called:
@ -147,11 +145,11 @@ class TestScript:
sc.start(tctx.options)
f = tflow.tflow(resp=True)
sc.request(f)
assert tctx.master.logs[0][0] == "error"
assert len(tctx.master.logs[0][1].splitlines()) == 6
assert re.search(r'addonscripts[\\/]error.py", line \d+, in request', tctx.master.logs[0][1])
assert re.search(r'addonscripts[\\/]error.py", line \d+, in mkerr', tctx.master.logs[0][1])
assert tctx.master.logs[0][1].endswith("ValueError: Error!\n")
assert tctx.master.logs[0].level == "error"
assert len(tctx.master.logs[0].msg.splitlines()) == 6
assert re.search(r'addonscripts[\\/]error.py", line \d+, in request', tctx.master.logs[0].msg)
assert re.search(r'addonscripts[\\/]error.py", line \d+, in mkerr', tctx.master.logs[0].msg)
assert tctx.master.logs[0].msg.endswith("ValueError: Error!\n")
def test_addon(self):
with taddons.context() as tctx:
@ -256,19 +254,19 @@ class TestScriptLoader:
"%s %s" % (rec, "c"),
]
)
debug = [(i[0], i[1]) for i in tctx.master.logs if i[0] == "debug"]
debug = [i.msg for i in tctx.master.logs if i.level == "debug"]
assert debug == [
('debug', 'a start'),
('debug', 'a configure'),
('debug', 'a running'),
'a start',
'a configure',
'a running',
('debug', 'b start'),
('debug', 'b configure'),
('debug', 'b running'),
'b start',
'b configure',
'b running',
('debug', 'c start'),
('debug', 'c configure'),
('debug', 'c running'),
'c start',
'c configure',
'c running',
]
tctx.master.logs = []
tctx.configure(
@ -279,8 +277,7 @@ class TestScriptLoader:
"%s %s" % (rec, "b"),
]
)
debug = [(i[0], i[1]) for i in tctx.master.logs if i[0] == "debug"]
# No events, only order has changed
debug = [i.msg for i in tctx.master.logs if i.level == "debug"]
assert debug == []
tctx.master.logs = []
@ -291,11 +288,11 @@ class TestScriptLoader:
"%s %s" % (rec, "a"),
]
)
debug = [(i[0], i[1]) for i in tctx.master.logs if i[0] == "debug"]
debug = [i.msg for i in tctx.master.logs if i.level == "debug"]
assert debug == [
('debug', 'c done'),
('debug', 'b done'),
('debug', 'x start'),
('debug', 'x configure'),
('debug', 'x running'),
'c done',
'b done',
'x start',
'x configure',
'x running',
]

View File

@ -250,17 +250,12 @@ class TestHTTP(tservers.HTTPProxyTest, CommonMixin):
assert p.request(req)
def test_get_connection_switching(self):
def switched(l):
for i in l:
if "serverdisconnect" in i:
return True
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 switched(self.proxy.tlog)
assert self.proxy.tmaster.has_log("serverdisconnect")
def test_blank_leading_line(self):
p = self.pathoc()
@ -602,7 +597,7 @@ 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 all("Error in handle_sni" not in msg for msg in self.proxy.tlog)
assert not self.proxy.tmaster.has_log("error in handle_sni")
def test_http(self):
p = self.pathoc(ssl=False)
@ -877,8 +872,7 @@ class TestServerConnect(tservers.HTTPProxyTest):
def test_unnecessary_serverconnect(self):
"""A replayed/fake response with no upstream_cert should not connect to an upstream server"""
assert self.pathod("200").status_code == 200
for msg in self.proxy.tmaster.tlog:
assert "serverconnect" not in msg
assert not self.proxy.tmaster.has_log("serverconnect")
class MasterKillRequest(tservers.TestMaster):

View File

@ -43,7 +43,7 @@ class TestConcurrent(tservers.MasterTest):
)
)
sc.start(tctx.options)
assert "decorator not supported" in tctx.master.logs[0][1]
assert tctx.master.has_log("decorator not supported")
def test_concurrent_class(self):
with taddons.context() as tctx:

View File

@ -7,6 +7,7 @@ from mitmproxy.exceptions import Kill, ControlException
from mitmproxy import controller
from mitmproxy import master
from mitmproxy import proxy
from mitmproxy.test import taddons
class TMsg:
@ -15,22 +16,18 @@ class TMsg:
class TestMaster:
def test_simple(self):
class DummyMaster(master.Master):
@controller.handler
class tAddon:
def log(self, _):
m.should_exit.set()
ctx.master.should_exit.set()
def tick(self, timeout):
# Speed up test
super().tick(0)
m = DummyMaster(None, proxy.DummyServer(None))
assert not m.should_exit.is_set()
msg = TMsg()
msg.reply = controller.DummyReply()
m.event_queue.put(("log", msg))
m.run()
assert m.should_exit.is_set()
with taddons.context() as ctx:
ctx.master.addons.add(tAddon())
assert not ctx.master.should_exit.is_set()
msg = TMsg()
msg.reply = controller.DummyReply()
ctx.master.event_queue.put(("log", msg))
ctx.master.run()
assert ctx.master.should_exit.is_set()
def test_server_simple(self):
m = master.Master(None, proxy.DummyServer(None))

View File

@ -6,7 +6,6 @@ import sys
import mitmproxy.platform
from mitmproxy.proxy.config import ProxyConfig
from mitmproxy.proxy.server import ProxyServer
from mitmproxy import master
from mitmproxy import controller
from mitmproxy import options
from mitmproxy import exceptions
@ -17,6 +16,7 @@ import pathod.pathoc
from mitmproxy.test import tflow
from mitmproxy.test import tutils
from mitmproxy.test import taddons
class MasterTest:
@ -68,11 +68,11 @@ class TestState:
# self.flows.append(f)
class TestMaster(master.Master):
class TestMaster(taddons.RecordingMaster):
def __init__(self, opts, config):
s = ProxyServer(config)
master.Master.__init__(self, opts, s)
super().__init__(opts, s)
def clear_addons(self, addons):
self.addons.clear()
@ -82,16 +82,9 @@ class TestMaster(master.Master):
self.addons.configure_all(self.options, self.options.keys())
self.addons.trigger("running")
def clear_log(self):
self.tlog = []
def reset(self, addons):
self.clear_addons(addons)
self.clear_log()
@controller.handler
def log(self, e):
self.tlog.append(e.msg)
self.clear()
class ProxyThread(threading.Thread):
@ -111,7 +104,7 @@ class ProxyThread(threading.Thread):
@property
def tlog(self):
return self.tmaster.tlog
return self.tmaster.logs
def run(self):
self.tmaster.run()