From a3da43d3e5d5b2ca243971e586aeee4969b6d053 Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Mon, 2 Apr 2018 14:51:14 +1200
Subject: 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.
---
 mitmproxy/addons/termstatus.py                  |  1 +
 mitmproxy/master.py                             |  8 ++++++++
 test/mitmproxy/data/addonscripts/shutdown.py    |  2 +-
 test/mitmproxy/proxy/protocol/test_websocket.py | 14 +++++++++-----
 test/mitmproxy/tools/test_main.py               | 10 ++++++++--
 5 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/mitmproxy/addons/termstatus.py b/mitmproxy/addons/termstatus.py
index c3c91283..3801f320 100644
--- a/mitmproxy/addons/termstatus.py
+++ b/mitmproxy/addons/termstatus.py
@@ -1,3 +1,4 @@
+import sys
 from mitmproxy import ctx
 from mitmproxy.utils import human
 
diff --git a/mitmproxy/master.py b/mitmproxy/master.py
index c22a78e4..022590d4 100644
--- a/mitmproxy/master.py
+++ b/mitmproxy/master.py
@@ -1,6 +1,7 @@
 import threading
 import contextlib
 import asyncio
+import logging
 
 from mitmproxy import addonmanager
 from mitmproxy import options
@@ -18,6 +19,13 @@ from mitmproxy.coretypes import basethread
 from . import ctx as mitmproxy_ctx
 
 
+# Conclusively preventing cross-thread races on proxy shutdown turns out to be
+# very hard. We could build a thread sync infrastructure for this, or we could
+# wait until we ditch threads and move all the protocols into the async loop.
+# Until then, silence non-critical errors.
+logging.getLogger('asyncio').setLevel(logging.CRITICAL)
+
+
 class ServerThread(basethread.BaseThread):
     def __init__(self, server):
         self.server = server
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
     ])
-- 
cgit v1.2.3