1 files changed, 66 insertions, 56 deletions
diff --git a/test/mitmproxy/test_addonmanager.py b/test/mitmproxy/test_addonmanager.py
index 5bb88eb6..8391e721 100644
--- a/test/mitmproxy/test_addonmanager.py
+++ b/test/mitmproxy/test_addonmanager.py
@@ -5,14 +5,17 @@ from mitmproxy import exceptions
 from mitmproxy import options
 from mitmproxy import master
 from mitmproxy import proxy
+from mitmproxy.test import taddons
 from mitmproxy.test import tflow
 
 
 class TAddon:
-    def __init__(self, name):
+    def __init__(self, name, addons=None):
         self.name = name
         self.tick = True
         self.custom_called = False
+        if addons:
+            self.addons = addons
 
     def __repr__(self):
         return "Addon(%s)" % self.name
@@ -34,19 +37,6 @@ class AOption:
         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))
@@ -70,40 +60,42 @@ def test_lifecycle():
     a = addonmanager.AddonManager(m)
     a.add(TAddon("one"))
 
+    with pytest.raises(exceptions.AddonError):
+        a.add(TAddon("one"))
+    with pytest.raises(exceptions.AddonError):
+        a.remove(TAddon("nonexistent"))
+
     f = tflow.tflow()
     a.handle_lifecycle("request", f)
 
-    a.configure_all(o, o.keys())
+    a._configure_all(o, o.keys())
 
 
 def test_simple():
-    o = options.Options()
-    m = master.Master(o, proxy.DummyServer(o))
-    a = addonmanager.AddonManager(m)
-    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"))
-    a.trigger("done")
-    with pytest.raises(exceptions.AddonError):
+    with taddons.context() as tctx:
+        a = tctx.master.addons
+
+        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"))
+        a.trigger("done")
         a.trigger("tick")
+        tctx.master.has_log("not callable")
 
-    a.remove(a.get("one"))
-    assert not a.get("one")
+        a.remove(a.get("one"))
+        assert not a.get("one")
 
-    ta = TAddon("one")
-    a.add(ta)
-    a.trigger("custom")
-    assert ta.custom_called
+        ta = TAddon("one")
+        a.add(ta)
+        a.trigger("custom")
+        assert ta.custom_called
 
 
 def test_load_option():
@@ -114,29 +106,47 @@ def test_load_option():
     assert "custom_option" in m.options._options
 
 
-def test_loadchain():
+def test_nesting():
     o = options.Options()
     m = master.Master(o, proxy.DummyServer(o))
     a = addonmanager.AddonManager(m)
 
-    a.add(AChain("one", None))
+    a.add(
+        TAddon(
+            "one",
+            addons=[
+                TAddon("two"),
+                TAddon("three", addons=[TAddon("four")])
+            ]
+        )
+    )
+    assert len(a.chain) == 1
     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()
+    assert a.get("four")
 
-    a.add(AChain("one", AChain("two", AChain("three", AChain("four", None)))))
-    assert not a.get("one")
-    assert not a.get("two")
+    a.trigger("custom")
+    assert a.get("one").custom_called
+    assert a.get("two").custom_called
+    assert a.get("three").custom_called
+    assert a.get("four").custom_called
+
+    a.remove(a.get("three"))
     assert not a.get("three")
-    assert a.get("four")
-    a.clear()
+    assert not a.get("four")
+
+
+class D:
+    def __init__(self):
+        self.w = None
+
+    def log(self, x):
+        self.w = x
+
+
+def test_streamlog():
+    dummy = D()
+    s = addonmanager.StreamLog(dummy.log)
+    s.write("foo")
+    assert dummy.w == "foo"