From 0c6663d0d5335e666598807c2e5d8f105803eac0 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Thu, 9 Mar 2017 13:52:58 +1300 Subject: Enable custom options for addons - Add an options parameter to the start() event. This is to be used by addons on startup to add custom options. - Add a running() event that is called once the proxy is up and running. - With the new paradigm we can't log during master __init__, so add a tiny termstatus addon to print proxy status to terminal once we're running. --- test/mitmproxy/addons/test_script.py | 46 ++++++++++++++-------- test/mitmproxy/addons/test_termstatus.py | 12 ++++++ test/mitmproxy/data/addonscripts/addon.py | 4 +- .../addonscripts/concurrent_decorator_class.py | 2 +- .../data/addonscripts/concurrent_decorator_err.py | 2 +- test/mitmproxy/data/addonscripts/recorder.py | 2 +- test/mitmproxy/proxy/test_server.py | 3 ++ test/mitmproxy/script/test_concurrent.py | 6 +-- test/mitmproxy/test_addonmanager.py | 8 ++-- test/mitmproxy/tools/console/test_master.py | 4 +- test/mitmproxy/tservers.py | 2 + 11 files changed, 62 insertions(+), 29 deletions(-) create mode 100644 test/mitmproxy/addons/test_termstatus.py (limited to 'test') diff --git a/test/mitmproxy/addons/test_script.py b/test/mitmproxy/addons/test_script.py index d79ed4ef..4c1b2e43 100644 --- a/test/mitmproxy/addons/test_script.py +++ b/test/mitmproxy/addons/test_script.py @@ -116,9 +116,7 @@ class TestScript: ) ) sc.load_script() - assert sc.ns.call_log == [ - ("solo", "start", (), {}), - ] + assert sc.ns.call_log[0][0:2] == ("solo", "start") sc.ns.call_log = [] f = tflow.tflow(resp=True) @@ -146,7 +144,7 @@ class TestScript: sc = script.Script( tutils.test_data.path("mitmproxy/data/addonscripts/error.py") ) - sc.start() + sc.start(tctx.options) f = tflow.tflow(resp=True) sc.request(f) assert tctx.master.event_log[0][0] == "error" @@ -162,7 +160,7 @@ class TestScript: "mitmproxy/data/addonscripts/addon.py" ) ) - sc.start() + sc.start(tctx.options) tctx.configure(sc) assert sc.ns.event_log == [ 'scriptstart', 'addonstart', 'addonconfigure' @@ -225,24 +223,31 @@ class TestScriptLoader: assert len(m.addons) == 1 def test_dupes(self): - o = options.Options(scripts=["one", "one"]) - m = master.Master(o, proxy.DummyServer()) sc = script.ScriptLoader() - with pytest.raises(exceptions.OptionsError): - m.addons.add(o, sc) + with taddons.context() as tctx: + tctx.master.addons.add(sc) + with pytest.raises(exceptions.OptionsError): + tctx.configure( + sc, + scripts = ["one", "one"] + ) def test_nonexistent(self): - o = options.Options(scripts=["nonexistent"]) - m = master.Master(o, proxy.DummyServer()) sc = script.ScriptLoader() - with pytest.raises(exceptions.OptionsError): - m.addons.add(o, sc) + with taddons.context() as tctx: + tctx.master.addons.add(sc) + with pytest.raises(exceptions.OptionsError): + tctx.configure( + sc, + scripts = ["nonexistent"] + ) def test_order(self): rec = tutils.test_data.path("mitmproxy/data/addonscripts/recorder.py") sc = script.ScriptLoader() with taddons.context() as tctx: tctx.master.addons.add(sc) + sc.running() tctx.configure( sc, scripts = [ @@ -253,9 +258,17 @@ class TestScriptLoader: ) debug = [(i[0], i[1]) for i in tctx.master.event_log if i[0] == "debug"] assert debug == [ - ('debug', 'a start'), ('debug', 'a configure'), - ('debug', 'b start'), ('debug', 'b configure'), - ('debug', 'c start'), ('debug', 'c configure') + ('debug', 'a start'), + ('debug', 'a configure'), + ('debug', 'a running'), + + ('debug', 'b start'), + ('debug', 'b configure'), + ('debug', 'b running'), + + ('debug', 'c start'), + ('debug', 'c configure'), + ('debug', 'c running'), ] tctx.master.event_log = [] tctx.configure( @@ -284,4 +297,5 @@ class TestScriptLoader: ('debug', 'b done'), ('debug', 'x start'), ('debug', 'x configure'), + ('debug', 'x running'), ] diff --git a/test/mitmproxy/addons/test_termstatus.py b/test/mitmproxy/addons/test_termstatus.py new file mode 100644 index 00000000..01c14814 --- /dev/null +++ b/test/mitmproxy/addons/test_termstatus.py @@ -0,0 +1,12 @@ +from mitmproxy.addons import termstatus +from mitmproxy.test import taddons + + +def test_configure(): + ts = termstatus.TermStatus() + with taddons.context() as ctx: + ts.running() + assert not ctx.master.event_log + ctx.configure(ts, server=True) + ts.running() + assert ctx.master.event_log diff --git a/test/mitmproxy/data/addonscripts/addon.py b/test/mitmproxy/data/addonscripts/addon.py index 84173cb6..f34f41cb 100644 --- a/test/mitmproxy/data/addonscripts/addon.py +++ b/test/mitmproxy/data/addonscripts/addon.py @@ -6,7 +6,7 @@ class Addon: def event_log(self): return event_log - def start(self): + def start(self, opts): event_log.append("addonstart") def configure(self, options, updated): @@ -17,6 +17,6 @@ def configure(options, updated): event_log.append("addonconfigure") -def start(): +def start(opts): event_log.append("scriptstart") return Addon() diff --git a/test/mitmproxy/data/addonscripts/concurrent_decorator_class.py b/test/mitmproxy/data/addonscripts/concurrent_decorator_class.py index bd047c99..10ba24cd 100644 --- a/test/mitmproxy/data/addonscripts/concurrent_decorator_class.py +++ b/test/mitmproxy/data/addonscripts/concurrent_decorator_class.py @@ -9,5 +9,5 @@ class ConcurrentClass: time.sleep(0.1) -def start(): +def start(opts): return ConcurrentClass() diff --git a/test/mitmproxy/data/addonscripts/concurrent_decorator_err.py b/test/mitmproxy/data/addonscripts/concurrent_decorator_err.py index 756869c8..7bc28182 100644 --- a/test/mitmproxy/data/addonscripts/concurrent_decorator_err.py +++ b/test/mitmproxy/data/addonscripts/concurrent_decorator_err.py @@ -2,5 +2,5 @@ from mitmproxy.script import concurrent @concurrent -def start(): +def start(opts): pass diff --git a/test/mitmproxy/data/addonscripts/recorder.py b/test/mitmproxy/data/addonscripts/recorder.py index 6b9b6ea8..aff524a8 100644 --- a/test/mitmproxy/data/addonscripts/recorder.py +++ b/test/mitmproxy/data/addonscripts/recorder.py @@ -22,5 +22,5 @@ class CallLogger: raise AttributeError -def start(): +def start(opts): return CallLogger(*sys.argv[1:]) diff --git a/test/mitmproxy/proxy/test_server.py b/test/mitmproxy/proxy/test_server.py index aa45761a..16efe415 100644 --- a/test/mitmproxy/proxy/test_server.py +++ b/test/mitmproxy/proxy/test_server.py @@ -302,6 +302,9 @@ class TestHTTP(tservers.HTTPProxyTest, CommonMixin): class TestHTTPAuth(tservers.HTTPProxyTest): def test_auth(self): self.master.addons.add(proxyauth.ProxyAuth()) + self.master.addons.configure_all( + self.master.options, self.master.options.keys() + ) self.master.options.proxyauth = "test:test" assert self.pathod("202").status_code == 407 p = self.pathoc() diff --git a/test/mitmproxy/script/test_concurrent.py b/test/mitmproxy/script/test_concurrent.py index e81c023d..a9b6f0c4 100644 --- a/test/mitmproxy/script/test_concurrent.py +++ b/test/mitmproxy/script/test_concurrent.py @@ -24,7 +24,7 @@ class TestConcurrent(tservers.MasterTest): "mitmproxy/data/addonscripts/concurrent_decorator.py" ) ) - sc.start() + sc.start(tctx.options) f1, f2 = tflow.tflow(), tflow.tflow() tctx.cycle(sc, f1) @@ -42,7 +42,7 @@ class TestConcurrent(tservers.MasterTest): "mitmproxy/data/addonscripts/concurrent_decorator_err.py" ) ) - sc.start() + sc.start(tctx.options) assert "decorator not supported" in tctx.master.event_log[0][1] def test_concurrent_class(self): @@ -52,7 +52,7 @@ class TestConcurrent(tservers.MasterTest): "mitmproxy/data/addonscripts/concurrent_decorator_class.py" ) ) - sc.start() + sc.start(tctx.options) f1, f2 = tflow.tflow(), tflow.tflow() tctx.cycle(sc, f1) diff --git a/test/mitmproxy/test_addonmanager.py b/test/mitmproxy/test_addonmanager.py index 17402e26..3e5f71c6 100644 --- a/test/mitmproxy/test_addonmanager.py +++ b/test/mitmproxy/test_addonmanager.py @@ -10,12 +10,12 @@ from mitmproxy import proxy class TAddon: def __init__(self, name): self.name = name - self.noop_member = True + self.tick = True def __repr__(self): return "Addon(%s)" % self.name - def noop(self): + def done(self): pass @@ -30,6 +30,6 @@ def test_simple(): assert not a.chain a.add(TAddon("one")) - a("noop") + a("done") with pytest.raises(exceptions.AddonError): - a("noop_member") + a("tick") diff --git a/test/mitmproxy/tools/console/test_master.py b/test/mitmproxy/tools/console/test_master.py index 0bf3734b..6c716ad1 100644 --- a/test/mitmproxy/tools/console/test_master.py +++ b/test/mitmproxy/tools/console/test_master.py @@ -28,7 +28,9 @@ class TestMaster(tservers.MasterTest): if "verbosity" not in opts: opts["verbosity"] = 1 o = options.Options(**opts) - return console.master.ConsoleMaster(o, proxy.DummyServer()) + m = console.master.ConsoleMaster(o, proxy.DummyServer()) + m.addons.configure_all(o, o.keys()) + return m def test_basic(self): m = self.mkmaster() diff --git a/test/mitmproxy/tservers.py b/test/mitmproxy/tservers.py index a8aaa358..c47411ee 100644 --- a/test/mitmproxy/tservers.py +++ b/test/mitmproxy/tservers.py @@ -79,6 +79,8 @@ class TestMaster(master.Master): self.state = TestState() self.addons.add(self.state) self.addons.add(*addons) + self.addons.configure_all(self.options, self.options.keys()) + self.addons.invoke_all_with_context("running") def clear_log(self): self.tlog = [] -- cgit v1.2.3