3 files changed, 57 insertions, 22 deletions
diff --git a/test/mitmproxy/builtins/test_script.py b/test/mitmproxy/builtins/test_script.py
index f37c7f94..826d2a91 100644
--- a/test/mitmproxy/builtins/test_script.py
+++ b/test/mitmproxy/builtins/test_script.py
@@ -48,7 +48,7 @@ def test_load_script():
             "data/addonscripts/recorder.py"
         ), []
     )
-    assert ns["configure"]
+    assert ns.start
 
 
 class TestScript(mastertest.MasterTest):
@@ -61,16 +61,16 @@ class TestScript(mastertest.MasterTest):
             )
         )
         m.addons.add(sc)
-        assert sc.ns["call_log"] == [
+        assert sc.ns.call_log == [
             ("solo", "start", (), {}),
             ("solo", "configure", (options.Options(),), {})
         ]
 
-        sc.ns["call_log"] = []
+        sc.ns.call_log = []
         f = tutils.tflow(resp=True)
         self.invoke(m, "request", f)
 
-        recf = sc.ns["call_log"][0]
+        recf = sc.ns.call_log[0]
         assert recf[1] == "request"
 
     def test_reload(self):
@@ -116,6 +116,19 @@ class TestScript(mastertest.MasterTest):
         assert not fm.state.view[0].request.is_replay
         assert fm.state.view[1].request.is_replay
 
+    def test_addon(self):
+        s = state.State()
+        m = master.FlowMaster(options.Options(), None, s)
+        sc = script.Script(
+            tutils.test_data.path(
+                "data/addonscripts/addon.py"
+            )
+        )
+        m.addons.add(sc)
+        assert sc.ns.event_log == [
+            'scriptstart', 'addonstart', 'addonconfigure'
+        ]
+
 
 class TestScriptLoader(mastertest.MasterTest):
     def test_simple(self):
diff --git a/test/mitmproxy/data/addonscripts/addon.py b/test/mitmproxy/data/addonscripts/addon.py
new file mode 100644
index 00000000..3b09d231
--- /dev/null
+++ b/test/mitmproxy/data/addonscripts/addon.py
@@ -0,0 +1,22 @@
+event_log = []
+
+
+class Addon:
+    @property
+    def event_log(self):
+        return event_log
+
+    def start(self):
+        event_log.append("addonstart")
+
+    def configure(self, options):
+        event_log.append("addonconfigure")
+
+
+def configure(options):
+    event_log.append("addonconfigure")
+
+
+def start():
+    event_log.append("scriptstart")
+    return Addon()
diff --git a/test/mitmproxy/data/addonscripts/recorder.py b/test/mitmproxy/data/addonscripts/recorder.py
index b6ac8d89..890e6f4e 100644
--- a/test/mitmproxy/data/addonscripts/recorder.py
+++ b/test/mitmproxy/data/addonscripts/recorder.py
@@ -2,24 +2,24 @@ from mitmproxy import controller
 from mitmproxy import ctx
 import sys
 
-call_log = []
 
-if len(sys.argv) > 1:
-    name = sys.argv[1]
-else:
-    name = "solo"
+class CallLogger:
+    call_log = []
 
-# Keep a log of all possible event calls
-evts = list(controller.Events) + ["configure"]
-for i in evts:
-    def mkprox():
-        evt = i
+    def __init__(self, name = "solo"):
+        self.name = name
 
-        def prox(*args, **kwargs):
-            lg = (name, evt, args, kwargs)
-            if evt != "log":
-                ctx.log.info(str(lg))
-            call_log.append(lg)
-            ctx.log.debug("%s %s" % (name, evt))
-        return prox
-    globals()[i] = mkprox()
+    def __getattr__(self, attr):
+        if attr in controller.Events:
+            def prox(*args, **kwargs):
+                lg = (self.name, attr, args, kwargs)
+                if attr != "log":
+                    ctx.log.info(str(lg))
+                self.call_log.append(lg)
+                ctx.log.debug("%s %s" % (self.name, attr))
+            return prox
+        raise AttributeError
+
+
+def start():
+    return CallLogger(*sys.argv[1:])