asyncio: test cleanup

Also silence asyncio logs. We sometimes end up with messages on the queue that
need to be ignored when the proxy shuts down, and asyncio complains loudly
about this.

3 files changed, 18 insertions, 8 deletions

diff --git a/test/mitmproxy/data/addonscripts/shutdown.py b/test/mitmproxy/data/addonscripts/shutdown.py
index 51a99b5c..3da4d03e 100644
--- a/test/mitmproxy/data/addonscripts/shutdown.py
+++ b/test/mitmproxy/data/addonscripts/shutdown.py
@@ -1,5 +1,5 @@
 from mitmproxy import ctx
 
 
-def running():
+def tick():
     ctx.master.shutdown()
diff --git a/test/mitmproxy/proxy/protocol/test_websocket.py b/test/mitmproxy/proxy/protocol/test_websocket.py
index 2a343450..014490b7 100644
--- a/test/mitmproxy/proxy/protocol/test_websocket.py
+++ b/test/mitmproxy/proxy/protocol/test_websocket.py
@@ -3,10 +3,10 @@ import os
 import struct
 import tempfile
 import traceback
+import time
 
 from mitmproxy import options
 from mitmproxy import exceptions
-from mitmproxy.addons import core
 from mitmproxy.http import HTTPFlow
 from mitmproxy.websocket import WebSocketFlow
 
@@ -54,6 +54,10 @@ class _WebSocketTestBase:
         cls.options = cls.get_options()
         cls.proxy = tservers.ProxyThread(tservers.TestMaster, cls.options)
         cls.proxy.start()
+        while True:
+            if cls.proxy.tmaster:
+                break
+            time.sleep(0.01)
 
     @classmethod
     def teardown_class(cls):
@@ -161,7 +165,7 @@ class TestSimple(_WebSocketTest):
             def websocket_start(self, f):
                 f.stream = streaming
 
-        self.master.addons.add(Stream())
+        self.proxy.set_addons(Stream())
         self.setup_connection()
 
         frame = websockets.Frame.from_file(self.client.rfile)
@@ -202,7 +206,7 @@ class TestSimple(_WebSocketTest):
             def websocket_message(self, f):
                 f.messages[-1].content = "foo"
 
-        self.master.addons.add(Addon())
+        self.proxy.set_addons(Addon())
         self.setup_connection()
 
         frame = websockets.Frame.from_file(self.client.rfile)
@@ -233,7 +237,7 @@ class TestKillFlow(_WebSocketTest):
             def websocket_message(self, f):
                 f.kill()
 
-        self.master.addons.add(KillFlow())
+        self.proxy.set_addons(KillFlow())
         self.setup_connection()
 
         with pytest.raises(exceptions.TcpDisconnect):
@@ -403,7 +407,7 @@ class TestStreaming(_WebSocketTest):
             def websocket_start(self, f):
                 f.stream = streaming
 
-        self.master.addons.add(Stream())
+        self.proxy.set_addons(Stream())
         self.setup_connection()
 
         frame = None
diff --git a/test/mitmproxy/tools/test_main.py b/test/mitmproxy/tools/test_main.py
index 88e2fe86..a514df74 100644
--- a/test/mitmproxy/tools/test_main.py
+++ b/test/mitmproxy/tools/test_main.py
@@ -1,19 +1,25 @@
+import pytest
+
 from mitmproxy.test import tutils
 from mitmproxy.tools import main
 
 shutdown_script = tutils.test_data.path("mitmproxy/data/addonscripts/shutdown.py")
 
 
-def test_mitmweb():
+@pytest.mark.asyncio
+async def test_mitmweb():
     main.mitmweb([
         "--no-web-open-browser",
         "-q",
+        "-p", "0",
         "-s", shutdown_script
     ])
 
 
-def test_mitmdump():
+@pytest.mark.asyncio
+async def test_mitmdump():
     main.mitmdump([
         "-q",
+        "-p", "0",
         "-s", shutdown_script
     ])