asyncio: shift script reloading out of the tick event

The tick event is a nasty compromise, left over from when we didn't have an
event loop. This is the first patch in a series that explores moving our
built-in addons to managing coroutines on the eventloop directly for periodic
tasks.

9 files changed, 80 insertions, 52 deletions

diff --git a/test/mitmproxy/addons/test_onboarding.py b/test/mitmproxy/addons/test_onboarding.py
index 0d99b1ff..a942062f 100644
--- a/test/mitmproxy/addons/test_onboarding.py
+++ b/test/mitmproxy/addons/test_onboarding.py
@@ -4,23 +4,21 @@ from mitmproxy.addons import onboarding
 from mitmproxy.test import taddons
 from .. import tservers
 
-import asyncio
-import tornado.platform.asyncio
-asyncio.set_event_loop_policy(tornado.platform.asyncio.AnyThreadEventLoopPolicy())
-
 
 class TestApp(tservers.HTTPProxyTest):
     def addons(self):
         return [onboarding.Onboarding()]
 
-    def test_basic(self):
+    @pytest.mark.asyncio
+    async def test_basic(self):
         ob = onboarding.Onboarding()
         with taddons.context(ob) as tctx:
             tctx.configure(ob)
             assert self.app("/").status_code == 200
 
     @pytest.mark.parametrize("ext", ["pem", "p12"])
-    def test_cert(self, ext):
+    @pytest.mark.asyncio
+    async def test_cert(self, ext):
         ob = onboarding.Onboarding()
         with taddons.context(ob) as tctx:
             tctx.configure(ob)
@@ -29,7 +27,8 @@ class TestApp(tservers.HTTPProxyTest):
             assert resp.content
 
     @pytest.mark.parametrize("ext", ["pem", "p12"])
-    def test_head(self, ext):
+    @pytest.mark.asyncio
+    async def test_head(self, ext):
         ob = onboarding.Onboarding()
         with taddons.context(ob) as tctx:
             tctx.configure(ob)
diff --git a/test/mitmproxy/addons/test_script.py b/test/mitmproxy/addons/test_script.py
index 96e19841..fad514ae 100644
--- a/test/mitmproxy/addons/test_script.py
+++ b/test/mitmproxy/addons/test_script.py
@@ -12,6 +12,10 @@ from mitmproxy.test import tflow
 from mitmproxy.test import tutils
 
 
+# We want this to be speedy for testing
+script.ReloadInterval = 0.1
+
+
 @pytest.mark.asyncio
 async def test_load_script():
     with taddons.context() as tctx:
@@ -71,7 +75,7 @@ class TestScript:
     def test_notfound(self):
         with taddons.context():
             with pytest.raises(exceptions.OptionsError):
-                script.Script("nonexistent")
+                script.Script("nonexistent", False)
 
     def test_quotes_around_filename(self):
         """
@@ -81,21 +85,23 @@ class TestScript:
         path = tutils.test_data.path("mitmproxy/data/addonscripts/recorder/recorder.py")
 
         s = script.Script(
-            '"{}"'.format(path)
+            '"{}"'.format(path),
+            False
         )
         assert '"' not in s.fullpath
 
-    def test_simple(self):
+    @pytest.mark.asyncio
+    async def test_simple(self):
         with taddons.context() as tctx:
             sc = script.Script(
                 tutils.test_data.path(
                     "mitmproxy/data/addonscripts/recorder/recorder.py"
-                )
+                ),
+                True,
             )
             tctx.master.addons.add(sc)
             tctx.configure(sc)
-            sc.tick()
-
+            await tctx.master.await_log("recorder running")
             rec = tctx.master.addons.get("recorder")
 
             assert rec.call_log[0][0:2] == ("recorder", "load")
@@ -112,25 +118,24 @@ class TestScript:
             f = tmpdir.join("foo.py")
             f.ensure(file=True)
             f.write("\n")
-            sc = script.Script(str(f))
+            sc = script.Script(str(f), True)
             tctx.configure(sc)
-            sc.tick()
             assert await tctx.master.await_log("Loading")
-            tctx.master.clear()
 
-            sc.last_load, sc.last_mtime = 0, 0
-            sc.tick()
+            tctx.master.clear()
+            f.write("\n")
             assert await tctx.master.await_log("Loading")
 
     @pytest.mark.asyncio
     async def test_exception(self):
         with taddons.context() as tctx:
             sc = script.Script(
-                tutils.test_data.path("mitmproxy/data/addonscripts/error.py")
+                tutils.test_data.path("mitmproxy/data/addonscripts/error.py"),
+                True,
             )
             tctx.master.addons.add(sc)
+            await tctx.master.await_log("error running")
             tctx.configure(sc)
-            sc.tick()
 
             f = tflow.tflow(resp=True)
             tctx.master.addons.trigger("request", f)
@@ -138,16 +143,17 @@ class TestScript:
             assert await tctx.master.await_log("ValueError: Error!")
             assert await tctx.master.await_log("error.py")
 
-    def test_addon(self):
+    @pytest.mark.asyncio
+    async def test_addon(self):
         with taddons.context() as tctx:
             sc = script.Script(
                 tutils.test_data.path(
                     "mitmproxy/data/addonscripts/addon.py"
-                )
+                ),
+                True
             )
             tctx.master.addons.add(sc)
-            tctx.configure(sc)
-            sc.tick()
+            await tctx.master.await_log("addon running")
             assert sc.ns.event_log == [
                 'scriptload', 'addonload', 'scriptconfigure', 'addonconfigure'
             ]
@@ -184,7 +190,6 @@ class TestScriptLoader:
             debug = [i.msg for i in tctx.master.logs if i.level == "debug"]
             assert debug == [
                 'recorder load', 'recorder running', 'recorder configure',
-                'recorder tick',
                 'recorder requestheaders', 'recorder request',
                 'recorder responseheaders', 'recorder response'
             ]
@@ -224,17 +229,21 @@ class TestScriptLoader:
                     scripts = ["one", "one"]
                 )
 
-    def test_script_deletion(self):
+    @pytest.mark.asyncio
+    async def test_script_deletion(self):
         tdir = tutils.test_data.path("mitmproxy/data/addonscripts/")
         with open(tdir + "/dummy.py", 'w') as f:
             f.write("\n")
+
         with taddons.context() as tctx:
             sl = script.ScriptLoader()
             tctx.master.addons.add(sl)
             tctx.configure(sl, scripts=[tutils.test_data.path("mitmproxy/data/addonscripts/dummy.py")])
+            await tctx.master.await_log("Loading")
 
             os.remove(tutils.test_data.path("mitmproxy/data/addonscripts/dummy.py"))
-            tctx.invoke(sl, "tick")
+
+            await tctx.master.await_log("Removing")
             assert not tctx.options.scripts
             assert not sl.addons
 
@@ -286,17 +295,14 @@ class TestScriptLoader:
                 'a load',
                 'a running',
                 'a configure',
-                'a tick',
 
                 'b load',
                 'b running',
                 'b configure',
-                'b tick',
 
                 'c load',
                 'c running',
                 'c configure',
-                'c tick',
             ]
 
             tctx.master.clear()
@@ -317,7 +323,7 @@ class TestScriptLoader:
                 'b configure',
             ]
 
-            tctx.master.logs = []
+            tctx.master.clear()
             tctx.configure(
                 sc,
                 scripts = [
@@ -325,9 +331,7 @@ class TestScriptLoader:
                     "%s/a.py" % rec,
                 ]
             )
-            tctx.master.addons.invoke_addon(sc, "tick")
-            await tctx.master.await_log("a tick")
-
+            await tctx.master.await_log("Loading")
             debug = [i.msg for i in tctx.master.logs if i.level == "debug"]
             assert debug == [
                 'c done',
@@ -336,6 +340,4 @@ class TestScriptLoader:
                 'e load',
                 'e running',
                 'e configure',
-                'e tick',
-                'a tick',
             ]
diff --git a/test/mitmproxy/data/addonscripts/addon.py b/test/mitmproxy/data/addonscripts/addon.py
index 8c834d82..c6b540d4 100644
--- a/test/mitmproxy/data/addonscripts/addon.py
+++ b/test/mitmproxy/data/addonscripts/addon.py
@@ -1,3 +1,4 @@
+from mitmproxy import ctx
 event_log = []
 
 
@@ -7,6 +8,7 @@ class Addon:
         return event_log
 
     def load(self, opts):
+        ctx.log.info("addon running")
         event_log.append("addonload")
 
     def configure(self, updated):
diff --git a/test/mitmproxy/data/addonscripts/error.py b/test/mitmproxy/data/addonscripts/error.py
index 4a3c370f..2f0c1755 100644
--- a/test/mitmproxy/data/addonscripts/error.py
+++ b/test/mitmproxy/data/addonscripts/error.py
@@ -1,6 +1,9 @@
-def mkerr():
-    raise ValueError("Error!")
+from mitmproxy import ctx
+
+
+def running():
+    ctx.log.info("error running")
 
 
 def request(flow):
-    mkerr()
+    raise ValueError("Error!")
diff --git a/test/mitmproxy/data/addonscripts/shutdown.py b/test/mitmproxy/data/addonscripts/shutdown.py
new file mode 100644
index 00000000..f4a8f55d
--- /dev/null
+++ b/test/mitmproxy/data/addonscripts/shutdown.py
@@ -0,0 +1,5 @@
+from mitmproxy import ctx
+
+
+def running():
+    ctx.master.shutdown()
\ No newline at end of file
diff --git a/test/mitmproxy/data/addonscripts/stream_modify.py b/test/mitmproxy/data/addonscripts/stream_modify.py
index 4fbf45c2..4ebcc3e9 100644
--- a/test/mitmproxy/data/addonscripts/stream_modify.py
+++ b/test/mitmproxy/data/addonscripts/stream_modify.py
@@ -1,7 +1,13 @@
+from mitmproxy import ctx
+
 def modify(chunks):
     for chunk in chunks:
         yield chunk.replace(b"foo", b"bar")
 
 
+def running():
+    ctx.log.info("stream_modify running")
+
+
 def responseheaders(flow):
     flow.response.stream = modify
diff --git a/test/mitmproxy/proxy/test_server.py b/test/mitmproxy/proxy/test_server.py
index bf24e28b..d7836104 100644
--- a/test/mitmproxy/proxy/test_server.py
+++ b/test/mitmproxy/proxy/test_server.py
@@ -256,11 +256,15 @@ class TestHTTP(tservers.HTTPProxyTest, CommonMixin):
             resp = p.request("get:'http://foo':h':foo'='bar'")
         assert resp.status_code == 400
 
-    def test_stream_modify(self):
+    @pytest.mark.asyncio
+    async def test_stream_modify(self):
         s = script.Script(
-            tutils.test_data.path("mitmproxy/data/addonscripts/stream_modify.py")
+            tutils.test_data.path("mitmproxy/data/addonscripts/stream_modify.py"),
+            False,
         )
         self.set_addons(s)
+        await self.master.await_log("stream_modify running")
+
         d = self.pathod('200:b"foo"')
         assert d.content == b"bar"
 
@@ -564,7 +568,8 @@ class TestTransparent(tservers.TransparentProxyTest, CommonMixin, TcpMixin):
 
     def test_tcp_stream_modify(self):
         s = script.Script(
-            tutils.test_data.path("mitmproxy/data/addonscripts/tcp_stream_modify.py")
+            tutils.test_data.path("mitmproxy/data/addonscripts/tcp_stream_modify.py"),
+            False,
         )
         self.set_addons(s)
         self._tcpproxy_on()
diff --git a/test/mitmproxy/tools/test_main.py b/test/mitmproxy/tools/test_main.py
index 57544276..751dcbe0 100644
--- a/test/mitmproxy/tools/test_main.py
+++ b/test/mitmproxy/tools/test_main.py
@@ -1,19 +1,23 @@
-import pytest
+import asyncio
 
 from mitmproxy.tools import main
-from mitmproxy import ctx
+from mitmproxy.test import tutils
 
+shutdown_script = tutils.test_data.path("mitmproxy/data/addonscripts/shutdown.py")
 
-@pytest.mark.asyncio
-async def test_mitmweb(event_loop):
+
+def test_mitmweb(event_loop):
+    asyncio.set_event_loop(event_loop)
     main.mitmweb([
         "--no-web-open-browser",
+        "-s", shutdown_script,
         "-q", "-p", "0",
     ])
-    await ctx.master._shutdown()
 
 
-@pytest.mark.asyncio
-async def test_mitmdump():
-    main.mitmdump(["-q", "-p", "0"])
-    await ctx.master._shutdown()
+def test_mitmdump(event_loop):
+    asyncio.set_event_loop(event_loop)
+    main.mitmdump([
+        "-s", shutdown_script,
+        "-q", "-p", "0",
+    ])
diff --git a/test/mitmproxy/tools/web/test_static_viewer.py b/test/mitmproxy/tools/web/test_static_viewer.py
index dfc45bc2..c044dee8 100644
--- a/test/mitmproxy/tools/web/test_static_viewer.py
+++ b/test/mitmproxy/tools/web/test_static_viewer.py
@@ -1,5 +1,6 @@
 import json
 from unittest import mock
+import pytest
 
 from mitmproxy.test import taddons
 from mitmproxy.test import tflow
@@ -57,7 +58,8 @@ def test_save_flows_content(ctx, tmpdir):
         assert p.join('response/content/Auto.json').check(file=1)
 
 
-def test_static_viewer(tmpdir):
+@pytest.mark.asyncio
+async def test_static_viewer(tmpdir):
     s = static_viewer.StaticViewer()
     rf = readfile.ReadFile()
     sa = save.Save()