asyncio: brutally rip out our old queue mechanism

6 files changed, 106 insertions, 97 deletions

diff --git a/test/mitmproxy/proxy/protocol/test_http2.py b/test/mitmproxy/proxy/protocol/test_http2.py
index d9aa03b4..8b929995 100644
--- a/test/mitmproxy/proxy/protocol/test_http2.py
+++ b/test/mitmproxy/proxy/protocol/test_http2.py
@@ -90,9 +90,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
diff --git a/test/mitmproxy/proxy/protocol/test_websocket.py b/test/mitmproxy/proxy/protocol/test_websocket.py
index 661605b7..2a343450 100644
--- a/test/mitmproxy/proxy/protocol/test_websocket.py
+++ b/test/mitmproxy/proxy/protocol/test_websocket.py
@@ -52,9 +52,7 @@ class _WebSocketTestBase:
     @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
diff --git a/test/mitmproxy/proxy/test_server.py b/test/mitmproxy/proxy/test_server.py
index 986dfb39..4cfaa523 100644
--- a/test/mitmproxy/proxy/test_server.py
+++ b/test/mitmproxy/proxy/test_server.py
@@ -21,14 +21,6 @@ from pathod import pathod
 from .. import tservers
 from ...conftest import skip_appveyor
 
-"""
-    Note that the choice of response code in these tests matters more than you
-    might think. libcurl treats a 304 response code differently from, say, a
-    200 response code - it will correctly terminate a 304 response with no
-    content-length header, whereas it will block forever waiting for content
-    for a 200 response.
-"""
-
 
 class CommonMixin:
 
diff --git a/test/mitmproxy/test_controller.py b/test/mitmproxy/test_controller.py
index e840380a..e27f6baf 100644
--- a/test/mitmproxy/test_controller.py
+++ b/test/mitmproxy/test_controller.py
@@ -1,7 +1,9 @@
+import asyncio
 from threading import Thread, Event
 from unittest.mock import Mock
 import queue
 import pytest
+import sys
 
 from mitmproxy.exceptions import Kill, ControlException
 from mitmproxy import controller
@@ -14,69 +16,87 @@ class TMsg:
     pass
 
 
-class TestMaster:
-    def test_simple(self):
-        class tAddon:
-            def log(self, _):
-                ctx.master.should_exit.set()
+def test_master():
+    class tAddon:
+        def log(self, _):
+            ctx.master.should_exit.set()
+
+    with taddons.context() as ctx:
+        ctx.master.addons.add(tAddon())
+        assert not ctx.master.should_exit.is_set()
 
-        with taddons.context() as ctx:
-            ctx.master.addons.add(tAddon())
-            assert not ctx.master.should_exit.is_set()
+        async def test():
             msg = TMsg()
             msg.reply = controller.DummyReply()
-            ctx.master.event_queue.put(("log", msg))
-            ctx.master.run()
-            assert ctx.master.should_exit.is_set()
-
-    def test_server_simple(self):
-        m = master.Master(None)
-        m.server = proxy.DummyServer()
-        m.start()
-        m.shutdown()
-        m.start()
-        m.shutdown()
-
-
-class TestServerThread:
-    def test_simple(self):
-        m = Mock()
-        t = master.ServerThread(m)
-        t.run()
-        assert m.serve_forever.called
-
-
-class TestChannel:
-    def test_tell(self):
-        q = queue.Queue()
-        channel = controller.Channel(q, Event())
-        m = Mock(name="test_tell")
-        channel.tell("test", m)
-        assert q.get() == ("test", m)
-        assert m.reply
-
-    def test_ask_simple(self):
-        q = queue.Queue()
-
-        def reply():
-            m, obj = q.get()
-            assert m == "test"
-            obj.reply.send(42)
-            obj.reply.take()
-            obj.reply.commit()
-
-        Thread(target=reply).start()
-
-        channel = controller.Channel(q, Event())
-        assert channel.ask("test", Mock(name="test_ask_simple")) == 42
-
-    def test_ask_shutdown(self):
-        q = queue.Queue()
-        done = Event()
-        done.set()
-        channel = controller.Channel(q, done)
-        with pytest.raises(Kill):
-            channel.ask("test", Mock(name="test_ask_shutdown"))
+            await ctx.master.event_queue.put(("log", msg))
+
+        ctx.master.run(inject=test)
+
+
+# class TestMaster:
+#     # def test_simple(self):
+#     #     class tAddon:
+#     #         def log(self, _):
+#     #             ctx.master.should_exit.set()
+
+#     #     with taddons.context() as ctx:
+#     #         ctx.master.addons.add(tAddon())
+#     #         assert not ctx.master.should_exit.is_set()
+#     #         msg = TMsg()
+#     #         msg.reply = controller.DummyReply()
+#     #         ctx.master.event_queue.put(("log", msg))
+#     #         ctx.master.run()
+#     #         assert ctx.master.should_exit.is_set()
+
+#     # def test_server_simple(self):
+#     #     m = master.Master(None)
+#     #     m.server = proxy.DummyServer()
+#     #     m.start()
+#     #     m.shutdown()
+#     #     m.start()
+#     #     m.shutdown()
+#     pass
+
+
+# class TestServerThread:
+#     def test_simple(self):
+#         m = Mock()
+#         t = master.ServerThread(m)
+#         t.run()
+#         assert m.serve_forever.called
+
+
+# class TestChannel:
+#     def test_tell(self):
+#         q = queue.Queue()
+#         channel = controller.Channel(q, Event())
+#         m = Mock(name="test_tell")
+#         channel.tell("test", m)
+#         assert q.get() == ("test", m)
+#         assert m.reply
+
+#     def test_ask_simple(self):
+#         q = queue.Queue()
+
+#         def reply():
+#             m, obj = q.get()
+#             assert m == "test"
+#             obj.reply.send(42)
+#             obj.reply.take()
+#             obj.reply.commit()
+
+#         Thread(target=reply).start()
+
+#         channel = controller.Channel(q, Event())
+#         assert channel.ask("test", Mock(name="test_ask_simple")) == 42
+
+#     def test_ask_shutdown(self):
+#         q = queue.Queue()
+#         done = Event()
+#         done.set()
+#         channel = controller.Channel(q, done)
+#         with pytest.raises(Kill):
+#             channel.ask("test", Mock(name="test_ask_shutdown"))
 
 
 class TestReply:
diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py
index 8cc11a16..9f1fb213 100644
--- a/test/mitmproxy/test_flow.py
+++ b/test/mitmproxy/test_flow.py
@@ -169,9 +169,10 @@ class TestFlowMaster:
         f.error = flow.Error("msg")
         fm.addons.handle_lifecycle("error", f)
 
-        fm.tell("foo", f)
-        with pytest.raises(ControlException):
-            fm.tick(timeout=1)
+        # FIXME: This no longer works, because we consume on the main loop.
+        # fm.tell("foo", f)
+        # with pytest.raises(ControlException):
+        #     fm.addons.trigger("unknown")
 
         fm.shutdown()
 
diff --git a/test/mitmproxy/tservers.py b/test/mitmproxy/tservers.py
index 0040b023..7be31a28 100644
--- a/test/mitmproxy/tservers.py
+++ b/test/mitmproxy/tservers.py
@@ -2,6 +2,7 @@ import os.path
 import threading
 import tempfile
 import sys
+import time
 from unittest import mock
 
 import mitmproxy.platform
@@ -62,11 +63,6 @@ class TestState:
         if f not in self.flows:
             self.flows.append(f)
 
-    # TODO: add TCP support?
-    # def tcp_start(self, f):
-    #     if f not in self.flows:
-    #         self.flows.append(f)
-
 
 class TestMaster(taddons.RecordingMaster):
 
@@ -90,13 +86,11 @@ class TestMaster(taddons.RecordingMaster):
 
 class ProxyThread(threading.Thread):
 
-    def __init__(self, tmaster):
+    def __init__(self, masterclass, options):
         threading.Thread.__init__(self)
-        self.tmaster = tmaster
-        self.name = "ProxyThread (%s:%s)" % (
-            tmaster.server.address[0],
-            tmaster.server.address[1],
-        )
+        self.masterclass = masterclass
+        self.options = options
+        self.tmaster = None
         controller.should_exit = False
 
     @property
@@ -107,12 +101,18 @@ class ProxyThread(threading.Thread):
     def tlog(self):
         return self.tmaster.logs
 
-    def run(self):
-        self.tmaster.run()
-
     def shutdown(self):
         self.tmaster.shutdown()
 
+    def run(self):
+        self.tmaster = self.masterclass(self.options)
+        self.tmaster.addons.add(core.Core())
+        self.name = "ProxyThread (%s:%s)" % (
+            self.tmaster.server.address[0],
+            self.tmaster.server.address[1],
+        )
+        self.tmaster.run()
+
 
 class ProxyTestBase:
     # Test Configuration
@@ -132,10 +132,12 @@ class ProxyTestBase:
             ssloptions=cls.ssloptions)
 
         cls.options = cls.get_options()
-        tmaster = cls.masterclass(cls.options)
-        tmaster.addons.add(core.Core())
-        cls.proxy = ProxyThread(tmaster)
+        cls.proxy = ProxyThread(cls.masterclass, cls.options)
         cls.proxy.start()
+        while True:
+            if cls.proxy.tmaster:
+                break
+            time.sleep(0.01)
 
     @classmethod
     def teardown_class(cls):
@@ -344,9 +346,7 @@ class ChainProxyTest(ProxyTestBase):
         cls.chain = []
         for _ in range(cls.n):
             opts = cls.get_options()
-            tmaster = cls.masterclass(opts)
-            tmaster.addons.add(core.Core())
-            proxy = ProxyThread(tmaster)
+            proxy = ProxyThread(cls.masterclass, opts)
             proxy.start()
             cls.chain.insert(0, proxy)