diff options
Diffstat (limited to 'test')
-rw-r--r-- | test/mitmproxy/addons/test_script.py | 9 | ||||
-rw-r--r-- | test/mitmproxy/data/addonscripts/addon.py | 4 | ||||
-rw-r--r-- | test/mitmproxy/data/addonscripts/concurrent_decorator_class.py | 4 | ||||
-rw-r--r-- | test/mitmproxy/data/addonscripts/recorder.py | 4 | ||||
-rw-r--r-- | test/mitmproxy/script/test_concurrent.py | 10 | ||||
-rw-r--r-- | test/mitmproxy/test_addonmanager.py | 95 |
6 files changed, 114 insertions, 12 deletions
diff --git a/test/mitmproxy/addons/test_script.py b/test/mitmproxy/addons/test_script.py index 79416161..c68981de 100644 --- a/test/mitmproxy/addons/test_script.py +++ b/test/mitmproxy/addons/test_script.py @@ -9,6 +9,7 @@ from unittest import mock from mitmproxy.test import tflow from mitmproxy.test import tutils from mitmproxy.test import taddons +from mitmproxy import addonmanager from mitmproxy import exceptions from mitmproxy import options from mitmproxy import proxy @@ -116,7 +117,7 @@ def test_script_print_stdout(): "mitmproxy/data/addonscripts/print.py" ), [] ) - ns.start(tctx.options) + ns.load(addonmanager.Loader(tctx.master)) mock_warn.assert_called_once_with("stdoutprint") @@ -157,7 +158,8 @@ class TestScript: sc = script.Script( tutils.test_data.path("mitmproxy/data/addonscripts/error.py") ) - sc.load(tctx.options) + l = addonmanager.Loader(tctx.master) + sc.load(l) f = tflow.tflow(resp=True) sc.request(f) assert tctx.master.logs[0].level == "error" @@ -173,7 +175,8 @@ class TestScript: "mitmproxy/data/addonscripts/addon.py" ) ) - sc.load(tctx.options) + l = addonmanager.Loader(tctx.master) + sc.load(l) tctx.configure(sc) assert sc.ns.event_log == [ 'scriptload', 'addonload', 'addonconfigure' diff --git a/test/mitmproxy/data/addonscripts/addon.py b/test/mitmproxy/data/addonscripts/addon.py index 3186400a..beef2ce7 100644 --- a/test/mitmproxy/data/addonscripts/addon.py +++ b/test/mitmproxy/data/addonscripts/addon.py @@ -17,6 +17,6 @@ def configure(options, updated): event_log.append("addonconfigure") -def load(opts): +def load(l): event_log.append("scriptload") - return Addon() + l.boot_into(Addon()) diff --git a/test/mitmproxy/data/addonscripts/concurrent_decorator_class.py b/test/mitmproxy/data/addonscripts/concurrent_decorator_class.py index d486923a..8e6988d4 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 load(opts): - return ConcurrentClass() +def load(l): + l.boot_into(ConcurrentClass()) diff --git a/test/mitmproxy/data/addonscripts/recorder.py b/test/mitmproxy/data/addonscripts/recorder.py index 5e0cdfc2..fe497b05 100644 --- a/test/mitmproxy/data/addonscripts/recorder.py +++ b/test/mitmproxy/data/addonscripts/recorder.py @@ -22,5 +22,5 @@ class CallLogger: raise AttributeError -def load(opts): - return CallLogger(*sys.argv[1:]) +def load(l): + l.boot_into(CallLogger(*sys.argv[1:])) diff --git a/test/mitmproxy/script/test_concurrent.py b/test/mitmproxy/script/test_concurrent.py index ccae0d31..86efdfc2 100644 --- a/test/mitmproxy/script/test_concurrent.py +++ b/test/mitmproxy/script/test_concurrent.py @@ -2,6 +2,7 @@ from mitmproxy.test import tflow from mitmproxy.test import tutils from mitmproxy.test import taddons +from mitmproxy import addonmanager from mitmproxy import controller from mitmproxy.addons import script @@ -24,7 +25,8 @@ class TestConcurrent(tservers.MasterTest): "mitmproxy/data/addonscripts/concurrent_decorator.py" ) ) - sc.load(tctx.options) + l = addonmanager.Loader(tctx.master) + sc.load(l) f1, f2 = tflow.tflow(), tflow.tflow() tctx.cycle(sc, f1) @@ -42,7 +44,8 @@ class TestConcurrent(tservers.MasterTest): "mitmproxy/data/addonscripts/concurrent_decorator_err.py" ) ) - sc.load(tctx.options) + l = addonmanager.Loader(tctx.master) + sc.load(l) assert tctx.master.has_log("decorator not supported") def test_concurrent_class(self): @@ -52,7 +55,8 @@ class TestConcurrent(tservers.MasterTest): "mitmproxy/data/addonscripts/concurrent_decorator_class.py" ) ) - sc.load(tctx.options) + l = addonmanager.Loader(tctx.master) + sc.load(l) 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 e7be25b8..5bb88eb6 100644 --- a/test/mitmproxy/test_addonmanager.py +++ b/test/mitmproxy/test_addonmanager.py @@ -5,6 +5,7 @@ from mitmproxy import exceptions from mitmproxy import options from mitmproxy import master from mitmproxy import proxy +from mitmproxy.test import tflow class TAddon: @@ -23,6 +24,58 @@ class TAddon: self.custom_called = True +class THalt: + def event_custom(self): + raise exceptions.AddonHalt + + +class AOption: + def load(self, l): + l.add_option("custom_option", bool, False, "help") + + +class AChain: + def __init__(self, name, next): + self.name = name + self.next = next + + def load(self, l): + if self.next: + l.boot_into(self.next) + + def __repr__(self): + return "<%s>" % self.name + + +def test_halt(): + o = options.Options() + m = master.Master(o, proxy.DummyServer(o)) + a = addonmanager.AddonManager(m) + halt = THalt() + end = TAddon("end") + a.add(halt) + a.add(end) + + a.trigger("custom") + assert not end.custom_called + + a.remove(halt) + a.trigger("custom") + assert end.custom_called + + +def test_lifecycle(): + o = options.Options() + m = master.Master(o, proxy.DummyServer(o)) + a = addonmanager.AddonManager(m) + a.add(TAddon("one")) + + f = tflow.tflow() + a.handle_lifecycle("request", f) + + a.configure_all(o, o.keys()) + + def test_simple(): o = options.Options() m = master.Master(o, proxy.DummyServer(o)) @@ -30,10 +83,13 @@ def test_simple(): with pytest.raises(exceptions.AddonError): a.invoke_addon(TAddon("one"), "done") + assert len(a) == 0 a.add(TAddon("one")) assert a.get("one") assert not a.get("two") + assert len(a) == 1 a.clear() + assert len(a) == 0 assert not a.chain a.add(TAddon("one")) @@ -41,7 +97,46 @@ def test_simple(): with pytest.raises(exceptions.AddonError): a.trigger("tick") + a.remove(a.get("one")) + assert not a.get("one") + ta = TAddon("one") a.add(ta) a.trigger("custom") assert ta.custom_called + + +def test_load_option(): + o = options.Options() + m = master.Master(o, proxy.DummyServer(o)) + a = addonmanager.AddonManager(m) + a.add(AOption()) + assert "custom_option" in m.options._options + + +def test_loadchain(): + o = options.Options() + m = master.Master(o, proxy.DummyServer(o)) + a = addonmanager.AddonManager(m) + + a.add(AChain("one", None)) + assert a.get("one") + a.clear() + + a.add(AChain("one", AChain("two", None))) + assert not a.get("one") + assert a.get("two") + a.clear() + + a.add(AChain("one", AChain("two", AChain("three", None)))) + assert not a.get("one") + assert not a.get("two") + assert a.get("three") + a.clear() + + a.add(AChain("one", AChain("two", AChain("three", AChain("four", None))))) + assert not a.get("one") + assert not a.get("two") + assert not a.get("three") + assert a.get("four") + a.clear() |