Use async for tctx.cycle/tctx.invoke.

mitmproxy/test/taddons.py

@@ -1,4 +1,3 @@
-import contextlib
 import asyncio
 import sys
 
@@ -79,15 +78,14 @@ class context:
     def __exit__(self, exc_type, exc_value, traceback):
         return False
 
-    @contextlib.contextmanager
+    async def cycle(self, addon, f):
-    def cycle(self, addon, f):
         """
             Cycles the flow through the events for the flow. Stops if a reply
             is taken (as in flow interception).
         """
         f.reply._state = "start"
         for evt in eventsequence.iterate(f):
-            self.master.addons.invoke_addon_sync(
+            await self.master.addons.invoke_addon(
                 addon,
                 evt
             )
@@ -115,11 +113,11 @@ class context:
         sc = script.Script(path, False)
         return sc.addons[0] if sc.addons else None
 
-    def invoke(self, addon, event: hooks.Hook):
+    async def invoke(self, addon, event: hooks.Hook):
         """
             Recursively invoke an event on an addon and all its children.
         """
-        return self.master.addons.invoke_addon_sync(addon, event)
+        return await self.master.addons.invoke_addon(addon, event)
 
     def command(self, func, *args):
         """

test/helper_tools/dumperview.py

@@ -1,4 +1,5 @@
 #!/usr/bin/env python3
+import asyncio
 import click
 
 from mitmproxy.addons import dumper
@@ -6,12 +7,24 @@ from mitmproxy.test import tflow
 from mitmproxy.test import taddons
 
 
+def run_async(coro):
+    """
+        Run the given async function in a new event loop.
+        This allows async functions to be called synchronously.
+    """
+    loop = asyncio.new_event_loop()
+    try:
+        return loop.run_until_complete(coro)
+    finally:
+        loop.close()
+
+
 def show(flow_detail, flows):
     d = dumper.Dumper()
     with taddons.context() as ctx:
         ctx.configure(d, flow_detail=flow_detail)
         for f in flows:
-            ctx.cycle(d, f)
+            run_async(ctx.cycle(d, f))
 
 
 @click.group()

test/mitmproxy/addons/test_intercept.py

@@ -7,7 +7,8 @@ from mitmproxy.test import taddons
 from mitmproxy.test import tflow
 
 
-def test_simple():
+@pytest.mark.asyncio
+async def test_simple():
     r = intercept.Intercept()
     with taddons.context(r) as tctx:
         assert not r.filt
@@ -23,11 +24,11 @@ def test_simple():
         tctx.configure(r, intercept="~s")
 
         f = tflow.tflow(resp=True)
-        tctx.cycle(r, f)
+        await tctx.cycle(r, f)
         assert f.intercepted
 
         f = tflow.tflow(resp=False)
-        tctx.cycle(r, f)
+        await tctx.cycle(r, f)
         assert not f.intercepted
 
         f = tflow.tflow(resp=True)
@@ -36,39 +37,41 @@ def test_simple():
 
         tctx.configure(r, intercept_active=False)
         f = tflow.tflow(resp=True)
-        tctx.cycle(r, f)
+        await tctx.cycle(r, f)
         assert not f.intercepted
 
         tctx.configure(r, intercept_active=True)
         f = tflow.tflow(resp=True)
-        tctx.cycle(r, f)
+        await tctx.cycle(r, f)
         assert f.intercepted
 
 
-def test_tcp():
+@pytest.mark.asyncio
+async def test_tcp():
     r = intercept.Intercept()
     with taddons.context(r) as tctx:
         tctx.configure(r, intercept="~tcp")
         f = tflow.ttcpflow()
-        tctx.cycle(r, f)
+        await tctx.cycle(r, f)
         assert f.intercepted
 
         tctx.configure(r, intercept_active=False)
         f = tflow.ttcpflow()
-        tctx.cycle(r, f)
+        await tctx.cycle(r, f)
         assert not f.intercepted
 
 
-def test_already_taken():
+@pytest.mark.asyncio
+async def test_already_taken():
     r = intercept.Intercept()
     with taddons.context(r) as tctx:
         tctx.configure(r, intercept="~q")
 
         f = tflow.tflow()
-        tctx.invoke(r, layers.http.HttpRequestHook(f))
+        await tctx.invoke(r, layers.http.HttpRequestHook(f))
         assert f.intercepted
 
         f = tflow.tflow()
         f.reply.take()
-        tctx.invoke(r, layers.http.HttpRequestHook(f))
+        await tctx.invoke(r, layers.http.HttpRequestHook(f))
         assert not f.intercepted

test/mitmproxy/addons/test_serverplayback.py

@@ -340,7 +340,8 @@ def test_server_playback_full():
         assert not tf.response
 
 
-def test_server_playback_kill():
+@pytest.mark.asyncio
+async def test_server_playback_kill():
     s = serverplayback.ServerPlayback()
     with taddons.context(s) as tctx:
         tctx.configure(
@@ -355,7 +356,7 @@ def test_server_playback_kill():
 
         f = tflow.tflow()
         f.request.host = "nonexistent"
-        tctx.cycle(s, f)
+        await tctx.cycle(s, f)
         assert f.error
 
 

test/mitmproxy/script/test_concurrent.py

@@ -1,3 +1,4 @@
+import asyncio
 import time
 
 import pytest
@@ -15,7 +16,8 @@ class Thing:
 
 
 class TestConcurrent:
-    def test_concurrent(self, tdata):
+    @pytest.mark.asyncio
+    async def test_concurrent(self, tdata):
         with taddons.context() as tctx:
             sc = tctx.script(
                 tdata.path(
@@ -23,8 +25,10 @@ class TestConcurrent:
                 )
             )
             f1, f2 = tflow.tflow(), tflow.tflow()
-            tctx.cycle(sc, f1)
+            await asyncio.gather(
-            tctx.cycle(sc, f2)
+                tctx.cycle(sc, f1),
+                tctx.cycle(sc, f2),
+            )
             start = time.time()
             while time.time() - start < 5:
                 if f1.reply.state == f2.reply.state == "committed":
@@ -41,7 +45,8 @@ class TestConcurrent:
             )
             await tctx.master.await_log("decorator not supported")
 
-    def test_concurrent_class(self, tdata):
+    @pytest.mark.asyncio
+    async def test_concurrent_class(self, tdata):
         with taddons.context() as tctx:
             sc = tctx.script(
                 tdata.path(
@@ -49,8 +54,10 @@ class TestConcurrent:
                 )
             )
             f1, f2 = tflow.tflow(), tflow.tflow()
-            tctx.cycle(sc, f1)
+            await asyncio.gather(
-            tctx.cycle(sc, f2)
+                tctx.cycle(sc, f1),
+                tctx.cycle(sc, f2),
+            )
             start = time.time()
             while time.time() - start < 5:
                 if f1.reply.state == f2.reply.state == "committed":