Merge pull request #2206 from cortesi/addonload

addon refactoring

8 files changed, 132 insertions, 30 deletions

diff --git a/test/mitmproxy/addons/test_script.py b/test/mitmproxy/addons/test_script.py
index 84d36e2a..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
@@ -104,7 +105,7 @@ def test_load_script():
                 "mitmproxy/data/addonscripts/recorder.py"
             ), []
         )
-        assert ns.start
+        assert ns.load
 
 
 def test_script_print_stdout():
@@ -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")
 
 
@@ -129,7 +130,7 @@ class TestScript:
                 )
             )
             sc.load_script()
-            assert sc.ns.call_log[0][0:2] == ("solo", "start")
+            assert sc.ns.call_log[0][0:2] == ("solo", "load")
 
             sc.ns.call_log = []
             f = tflow.tflow(resp=True)
@@ -157,7 +158,8 @@ class TestScript:
             sc = script.Script(
                 tutils.test_data.path("mitmproxy/data/addonscripts/error.py")
             )
-            sc.start(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,10 +175,11 @@ class TestScript:
                     "mitmproxy/data/addonscripts/addon.py"
                 )
             )
-            sc.start(tctx.options)
+            l = addonmanager.Loader(tctx.master)
+            sc.load(l)
             tctx.configure(sc)
             assert sc.ns.event_log == [
-                'scriptstart', 'addonstart', 'addonconfigure'
+                'scriptload', 'addonload', 'addonconfigure'
             ]
 
 
@@ -213,7 +216,7 @@ class TestScriptLoader:
                 ), [f]
             )
         evts = [i[1] for i in sc.ns.call_log]
-        assert evts == ['start', 'requestheaders', 'request', 'responseheaders', 'response', 'done']
+        assert evts == ['load', 'requestheaders', 'request', 'responseheaders', 'response', 'done']
 
         f = tflow.tflow(resp=True)
         with m.handlecontext():
@@ -271,15 +274,15 @@ class TestScriptLoader:
             )
             debug = [i.msg for i in tctx.master.logs if i.level == "debug"]
             assert debug == [
-                'a start',
+                'a load',
                 'a configure',
                 'a running',
 
-                'b start',
+                'b load',
                 'b configure',
                 'b running',
 
-                'c start',
+                'c load',
                 'c configure',
                 'c running',
             ]
@@ -307,7 +310,7 @@ class TestScriptLoader:
             assert debug == [
                 'c done',
                 'b done',
-                'x start',
+                'x load',
                 'x configure',
                 'x running',
             ]
diff --git a/test/mitmproxy/addons/test_serverplayback.py b/test/mitmproxy/addons/test_serverplayback.py
index 02642c35..7078b66e 100644
--- a/test/mitmproxy/addons/test_serverplayback.py
+++ b/test/mitmproxy/addons/test_serverplayback.py
@@ -44,12 +44,12 @@ def test_server_playback():
 
     assert not sp.flowmap
 
-    sp.load([f])
+    sp.load_flows([f])
     assert sp.flowmap
     assert sp.next_flow(f)
     assert not sp.flowmap
 
-    sp.load([f])
+    sp.load_flows([f])
     assert sp.flowmap
     sp.clear()
     assert not sp.flowmap
@@ -192,7 +192,7 @@ def test_load():
     r2 = tflow.tflow(resp=True)
     r2.request.headers["key"] = "two"
 
-    s.load([r, r2])
+    s.load_flows([r, r2])
 
     assert s.count() == 2
 
@@ -218,7 +218,7 @@ def test_load_with_server_replay_nopop():
     r2 = tflow.tflow(resp=True)
     r2.request.headers["key"] = "two"
 
-    s.load([r, r2])
+    s.load_flows([r, r2])
 
     assert s.count() == 2
     s.next_flow(r)
@@ -319,7 +319,7 @@ def test_server_playback_full():
 
         f = tflow.tflow()
         f.response = mitmproxy.test.tutils.tresp(content=f.request.content)
-        s.load([f, f])
+        s.load_flows([f, f])
 
         tf = tflow.tflow()
         assert not tf.response
@@ -352,7 +352,7 @@ def test_server_playback_kill():
 
         f = tflow.tflow()
         f.response = mitmproxy.test.tutils.tresp(content=f.request.content)
-        s.load([f])
+        s.load_flows([f])
 
         f = tflow.tflow()
         f.request.host = "nonexistent"
diff --git a/test/mitmproxy/data/addonscripts/addon.py b/test/mitmproxy/data/addonscripts/addon.py
index f34f41cb..beef2ce7 100644
--- a/test/mitmproxy/data/addonscripts/addon.py
+++ b/test/mitmproxy/data/addonscripts/addon.py
@@ -6,8 +6,8 @@ class Addon:
     def event_log(self):
         return event_log
 
-    def start(self, opts):
-        event_log.append("addonstart")
+    def load(self, opts):
+        event_log.append("addonload")
 
     def configure(self, options, updated):
         event_log.append("addonconfigure")
@@ -17,6 +17,6 @@ def configure(options, updated):
     event_log.append("addonconfigure")
 
 
-def start(opts):
-    event_log.append("scriptstart")
-    return Addon()
+def load(l):
+    event_log.append("scriptload")
+    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 10ba24cd..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 start(opts):
-    return ConcurrentClass()
+def load(l):
+    l.boot_into(ConcurrentClass())
diff --git a/test/mitmproxy/data/addonscripts/print.py b/test/mitmproxy/data/addonscripts/print.py
index fdfcc3d9..93b65a64 100644
--- a/test/mitmproxy/data/addonscripts/print.py
+++ b/test/mitmproxy/data/addonscripts/print.py
@@ -1,2 +1,2 @@
-def start(opts):
+def load(l):
     print("stdoutprint")
diff --git a/test/mitmproxy/data/addonscripts/recorder.py b/test/mitmproxy/data/addonscripts/recorder.py
index aff524a8..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 start(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 0e397b8f..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.start(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.start(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.start(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()