From 9294d19f908cdde86b390fb2fa93aad422da8522 Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Sat, 23 Jul 2016 11:28:44 +1200
Subject: [PATCH] scripts: .start can now return an Addon instance

---
 mitmproxy/builtins/script.py              |  6 ++++--
 test/mitmproxy/builtins/test_script.py    | 13 +++++++++++++
 test/mitmproxy/data/addonscripts/addon.py | 22 ++++++++++++++++++++++
 3 files changed, 39 insertions(+), 2 deletions(-)
 create mode 100644 test/mitmproxy/data/addonscripts/addon.py

diff --git a/mitmproxy/builtins/script.py b/mitmproxy/builtins/script.py
index 220fb9af7..ed5f2ecd6 100644
--- a/mitmproxy/builtins/script.py
+++ b/mitmproxy/builtins/script.py
@@ -99,7 +99,6 @@ class ReloadHandler(watchdog.events.FileSystemEventHandler):
             return False
         if os.path.basename(event.src_path).startswith("."):
             return False
-        print(event.src_path)
         return True
 
     def on_modified(self, event):
@@ -151,7 +150,10 @@ class Script:
 
     def load_script(self):
         self.ns = load_script(self.path, self.args)
-        self.run("start")
+        ret = self.run("start")
+        if ret:
+            self.ns = ret
+            self.run("start")
 
     def tick(self):
         if self.should_reload.is_set():
diff --git a/test/mitmproxy/builtins/test_script.py b/test/mitmproxy/builtins/test_script.py
index 1658088c7..3f07b5767 100644
--- a/test/mitmproxy/builtins/test_script.py
+++ b/test/mitmproxy/builtins/test_script.py
@@ -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 000000000..3b09d231a
--- /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()