3 files changed, 16 insertions, 19 deletions
diff --git a/test/mitmproxy/proxy/protocol/test_http2.py b/test/mitmproxy/proxy/protocol/test_http2.py
index d9aa03b4..b5f21413 100644
--- a/test/mitmproxy/proxy/protocol/test_http2.py
+++ b/test/mitmproxy/proxy/protocol/test_http2.py
@@ -10,7 +10,6 @@ import h2
 from mitmproxy import options
 
 import mitmproxy.net
-from mitmproxy.addons import core
 from ...net import tservers as net_tservers
 from mitmproxy import exceptions
 from mitmproxy.net.http import http1, http2
@@ -90,9 +89,7 @@ class _Http2TestBase:
     @classmethod
     def setup_class(cls):
         cls.options = cls.get_options()
-        tmaster = tservers.TestMaster(cls.options)
-        tmaster.addons.add(core.Core())
-        cls.proxy = tservers.ProxyThread(tmaster)
+        cls.proxy = tservers.ProxyThread(tservers.TestMaster, cls.options)
         cls.proxy.start()
 
     @classmethod
@@ -106,7 +103,7 @@ class _Http2TestBase:
             upstream_cert=True,
             ssl_insecure=True
         )
-        opts.cadir = os.path.join(tempfile.gettempdir(), "mitmproxy")
+        opts.confdir = os.path.join(tempfile.gettempdir(), "mitmproxy")
         return opts
 
     @property
@@ -120,6 +117,7 @@ class _Http2TestBase:
     def teardown(self):
         if self.client:
             self.client.close()
+        self.server.server.wait_for_silence()
 
     def setup_connection(self):
         self.client = mitmproxy.net.tcp.TCPClient(("127.0.0.1", self.proxy.port))
diff --git a/test/mitmproxy/proxy/protocol/test_http_replay.py b/test/mitmproxy/proxy/protocol/test_http_replay.py
deleted file mode 100644
index 777ab4dd..00000000
--- a/test/mitmproxy/proxy/protocol/test_http_replay.py
+++ /dev/null
@@ -1 +0,0 @@
-# TODO: write tests
diff --git a/test/mitmproxy/proxy/protocol/test_websocket.py b/test/mitmproxy/proxy/protocol/test_websocket.py
index 661605b7..1f4e2bca 100644
--- a/test/mitmproxy/proxy/protocol/test_websocket.py
+++ b/test/mitmproxy/proxy/protocol/test_websocket.py
@@ -6,7 +6,6 @@ import traceback
 
 from mitmproxy import options
 from mitmproxy import exceptions
-from mitmproxy.addons import core
 from mitmproxy.http import HTTPFlow
 from mitmproxy.websocket import WebSocketFlow
 
@@ -48,13 +47,12 @@ class _WebSocketServerBase(net_tservers.ServerTestBase):
 
 
 class _WebSocketTestBase:
+    client = None
 
     @classmethod
     def setup_class(cls):
         cls.options = cls.get_options()
-        tmaster = tservers.TestMaster(cls.options)
-        tmaster.addons.add(core.Core())
-        cls.proxy = tservers.ProxyThread(tmaster)
+        cls.proxy = tservers.ProxyThread(tservers.TestMaster, cls.options)
         cls.proxy.start()
 
     @classmethod
@@ -69,7 +67,7 @@ class _WebSocketTestBase:
             ssl_insecure=True,
             websocket=True,
         )
-        opts.cadir = os.path.join(tempfile.gettempdir(), "mitmproxy")
+        opts.confdir = os.path.join(tempfile.gettempdir(), "mitmproxy")
         return opts
 
     @property
@@ -163,7 +161,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)
@@ -204,7 +202,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)
@@ -235,7 +233,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):
@@ -288,7 +286,8 @@ class TestPing(_WebSocketTest):
         wfile.flush()
         websockets.Frame.from_file(rfile)
 
-    def test_ping(self):
+    @pytest.mark.asyncio
+    async def test_ping(self):
         self.setup_connection()
 
         frame = websockets.Frame.from_file(self.client.rfile)
@@ -298,7 +297,7 @@ class TestPing(_WebSocketTest):
         assert frame.header.opcode == websockets.OPCODE.PING
         assert frame.payload == b''  # We don't send payload to other end
 
-        assert self.master.has_log("Pong Received from server", "info")
+        assert await self.master.await_log("Pong Received from server", "info")
 
 
 class TestPong(_WebSocketTest):
@@ -316,7 +315,8 @@ class TestPong(_WebSocketTest):
         wfile.flush()
         websockets.Frame.from_file(rfile)
 
-    def test_pong(self):
+    @pytest.mark.asyncio
+    async def test_pong(self):
         self.setup_connection()
 
         self.client.wfile.write(bytes(websockets.Frame(fin=1, mask=1, opcode=websockets.OPCODE.PING, payload=b'foobar')))
@@ -329,7 +329,7 @@ class TestPong(_WebSocketTest):
 
         assert frame.header.opcode == websockets.OPCODE.PONG
         assert frame.payload == b'foobar'
-        assert self.master.has_log("Pong Received from server", "info")
+        assert await self.master.await_log("pong received")
 
 
 class TestClose(_WebSocketTest):
@@ -405,7 +405,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