diff options
| -rw-r--r-- | mitmproxy/master.py | 18 | ||||
| -rw-r--r-- | mitmproxy/tools/main.py | 5 | ||||
| -rw-r--r-- | test/mitmproxy/addons/test_keepserving.py | 7 | ||||
| -rw-r--r-- | test/mitmproxy/data/addonscripts/shutdown.py | 5 | ||||
| -rw-r--r-- | test/mitmproxy/tools/test_main.py | 19 | ||||
| -rw-r--r-- | test/mitmproxy/tservers.py | 2 |
6 files changed, 29 insertions, 27 deletions
diff --git a/mitmproxy/master.py b/mitmproxy/master.py index dff8ca5c..19a2ac67 100644 --- a/mitmproxy/master.py +++ b/mitmproxy/master.py @@ -92,19 +92,27 @@ class Master: try: loop.run_forever() finally: - self.shutdown() pending = asyncio.Task.all_tasks() loop.run_until_complete(asyncio.gather(*pending)) loop.close() self.addons.trigger("done") + async def _shutdown(self): + if self.server: + self.server.shutdown() + loop = asyncio.get_event_loop() + loop.stop() + def shutdown(self): + """ + Shut down the proxy. This method is thread-safe. + """ if not self.should_exit.is_set(): - if self.server: - self.server.shutdown() self.should_exit.set() - loop = asyncio.get_event_loop() - loop.stop() + asyncio.run_coroutine_threadsafe( + self._shutdown(), + loop = self.channel.loop, + ) def _change_reverse_host(self, f): """ diff --git a/mitmproxy/tools/main.py b/mitmproxy/tools/main.py index e8ec993d..8d63aa32 100644 --- a/mitmproxy/tools/main.py +++ b/mitmproxy/tools/main.py @@ -140,7 +140,7 @@ def mitmproxy(args=None): # pragma: no cover assert_utf8_env() from mitmproxy.tools import console - run(console.master.ConsoleMaster, cmdline.mitmproxy, args) + return run(console.master.ConsoleMaster, cmdline.mitmproxy, args) def mitmdump(args=None): # pragma: no cover @@ -159,8 +159,9 @@ def mitmdump(args=None): # pragma: no cover m = run(dump.DumpMaster, cmdline.mitmdump, args, extra) if m and m.errorcheck.has_errored: sys.exit(1) + return m def mitmweb(args=None): # pragma: no cover from mitmproxy.tools import web - run(web.master.WebMaster, cmdline.mitmweb, args) + return run(web.master.WebMaster, cmdline.mitmweb, args) diff --git a/test/mitmproxy/addons/test_keepserving.py b/test/mitmproxy/addons/test_keepserving.py index 2869d097..5eafa792 100644 --- a/test/mitmproxy/addons/test_keepserving.py +++ b/test/mitmproxy/addons/test_keepserving.py @@ -1,9 +1,14 @@ +import asyncio +import pytest + from mitmproxy.addons import keepserving from mitmproxy.test import taddons -def test_keepserving(): +@pytest.mark.asyncio +async def test_keepserving(): ks = keepserving.KeepServing() with taddons.context(ks) as tctx: ks.event_processing_complete() + asyncio.sleep(0.1) assert tctx.master.should_exit.is_set() diff --git a/test/mitmproxy/data/addonscripts/shutdown.py b/test/mitmproxy/data/addonscripts/shutdown.py deleted file mode 100644 index 3da4d03e..00000000 --- a/test/mitmproxy/data/addonscripts/shutdown.py +++ /dev/null @@ -1,5 +0,0 @@ -from mitmproxy import ctx - - -def tick(): - ctx.master.shutdown() diff --git a/test/mitmproxy/tools/test_main.py b/test/mitmproxy/tools/test_main.py index a514df74..ba213733 100644 --- a/test/mitmproxy/tools/test_main.py +++ b/test/mitmproxy/tools/test_main.py @@ -1,25 +1,18 @@ import pytest -from mitmproxy.test import tutils from mitmproxy.tools import main -shutdown_script = tutils.test_data.path("mitmproxy/data/addonscripts/shutdown.py") - @pytest.mark.asyncio -async def test_mitmweb(): - main.mitmweb([ +async def test_mitmweb(event_loop): + m = main.mitmweb([ "--no-web-open-browser", - "-q", - "-p", "0", - "-s", shutdown_script + "-q", "-p", "0", ]) + await m._shutdown() @pytest.mark.asyncio async def test_mitmdump(): - main.mitmdump([ - "-q", - "-p", "0", - "-s", shutdown_script - ]) + m = main.mitmdump(["-q", "-p", "0"]) + await m._shutdown() diff --git a/test/mitmproxy/tservers.py b/test/mitmproxy/tservers.py index 13ee4a43..857ca45d 100644 --- a/test/mitmproxy/tservers.py +++ b/test/mitmproxy/tservers.py @@ -40,7 +40,7 @@ class MasterTest: async def dummy_cycle(self, master, n, content): for i in range(n): await self.cycle(master, content) - master.shutdown() + await master._shutdown() def flowfile(self, path): with open(path, "wb") as f: |