From caf49300c21fe05a5ee521270f0ea96c5f10c652 Mon Sep 17 00:00:00 2001 From: Robert Xiao Date: Thu, 3 Feb 2022 03:25:11 -0800 Subject: [PATCH] Use async for tctx.cycle/tctx.invoke. --- mitmproxy/test/taddons.py | 10 ++++---- test/helper_tools/dumperview.py | 15 +++++++++++- test/mitmproxy/addons/test_intercept.py | 25 +++++++++++--------- test/mitmproxy/addons/test_serverplayback.py | 5 ++-- test/mitmproxy/script/test_concurrent.py | 19 ++++++++++----- 5 files changed, 48 insertions(+), 26 deletions(-) diff --git a/mitmproxy/test/taddons.py b/mitmproxy/test/taddons.py index ec6b23b8f..1129072a9 100644 --- a/mitmproxy/test/taddons.py +++ b/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 - def cycle(self, addon, f): + async 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): """ diff --git a/test/helper_tools/dumperview.py b/test/helper_tools/dumperview.py index b6ade8c2b..4b12ceb57 100755 --- a/test/helper_tools/dumperview.py +++ b/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() diff --git a/test/mitmproxy/addons/test_intercept.py b/test/mitmproxy/addons/test_intercept.py index aad4e2c6f..615dfa906 100644 --- a/test/mitmproxy/addons/test_intercept.py +++ b/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 diff --git a/test/mitmproxy/addons/test_serverplayback.py b/test/mitmproxy/addons/test_serverplayback.py index 9e0ca1566..19d5ec53d 100644 --- a/test/mitmproxy/addons/test_serverplayback.py +++ b/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 diff --git a/test/mitmproxy/script/test_concurrent.py b/test/mitmproxy/script/test_concurrent.py index 9743b7344..8b75967d6 100644 --- a/test/mitmproxy/script/test_concurrent.py +++ b/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) - tctx.cycle(sc, f2) + await asyncio.gather( + 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) - tctx.cycle(sc, f2) + await asyncio.gather( + tctx.cycle(sc, f1), + tctx.cycle(sc, f2), + ) start = time.time() while time.time() - start < 5: if f1.reply.state == f2.reply.state == "committed":