cleanup mitmproxy.controller, raise Kill in Channel (#1085)

4 files changed, 109 insertions, 18 deletions

diff --git a/test/mitmproxy/test_controller.py b/test/mitmproxy/test_controller.py
index eb3f7df4..f7bf615a 100644
--- a/test/mitmproxy/test_controller.py
+++ b/test/mitmproxy/test_controller.py
@@ -1,11 +1,105 @@
-import mock
-from mitmproxy import controller
+from threading import Thread, Event
 
+from mock import Mock
 
-class TestMaster:
+from mitmproxy.controller import Reply, DummyReply, Channel, ServerThread, ServerMaster, Master
+from six.moves import queue
 
-    def test_default_handler(self):
-        m = controller.Master(None)
-        msg = mock.MagicMock()
-        m.handle("type", msg)
-        assert msg.reply.call_count == 1
+from mitmproxy.exceptions import Kill
+from mitmproxy.proxy import DummyServer
+from netlib.tutils import raises
+
+
+class TestMaster(object):
+    def test_simple(self):
+
+        class DummyMaster(Master):
+            def handle_panic(self, _):
+                m.should_exit.set()
+
+            def tick(self, timeout):
+                # Speed up test
+                super(DummyMaster, self).tick(0)
+
+        m = DummyMaster()
+        assert not m.should_exit.is_set()
+        m.event_queue.put(("panic", 42))
+        m.run()
+        assert m.should_exit.is_set()
+
+
+class TestServerMaster(object):
+    def test_simple(self):
+        m = ServerMaster()
+        s = DummyServer(None)
+        m.add_server(s)
+        m.start()
+        m.shutdown()
+        m.start()
+        m.shutdown()
+
+
+class TestServerThread(object):
+    def test_simple(self):
+        m = Mock()
+        t = ServerThread(m)
+        t.run()
+        assert m.serve_forever.called
+
+
+class TestChannel(object):
+    def test_tell(self):
+        q = queue.Queue()
+        channel = Channel(q, Event())
+        m = Mock()
+        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(42)
+
+        Thread(target=reply).start()
+
+        channel = Channel(q, Event())
+        assert channel.ask("test", Mock()) == 42
+
+    def test_ask_shutdown(self):
+        q = queue.Queue()
+        done = Event()
+        done.set()
+        channel = Channel(q, done)
+        with raises(Kill):
+            channel.ask("test", Mock())
+
+
+class TestDummyReply(object):
+    def test_simple(self):
+        reply = DummyReply()
+        assert not reply.acked
+        reply()
+        assert reply.acked
+
+
+class TestReply(object):
+    def test_simple(self):
+        reply = Reply(42)
+        assert not reply.acked
+        reply("foo")
+        assert reply.acked
+        assert reply.q.get() == "foo"
+
+    def test_default(self):
+        reply = Reply(42)
+        reply()
+        assert reply.q.get() == 42
+
+    def test_reply_none(self):
+        reply = Reply(42)
+        reply(None)
+        assert reply.q.get() is None
diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py
index 7ede7a81..60f6b1a9 100644
--- a/test/mitmproxy/test_flow.py
+++ b/test/mitmproxy/test_flow.py
@@ -116,9 +116,8 @@ class TestClientPlaybackState:
         c.clear(c.current)
         assert c.done()
 
-        q = queue.Queue()
         fm.state.clear()
-        fm.tick(q, timeout=0)
+        fm.tick(timeout=0)
 
         fm.stop_client_playback()
         assert not fm.client_playback
@@ -858,9 +857,8 @@ class TestFlowMaster:
         assert not fm.start_client_playback(pb, False)
         fm.client_playback.testing = True
 
-        q = queue.Queue()
         assert not fm.state.flow_count()
-        fm.tick(q, 0)
+        fm.tick(0)
         assert fm.state.flow_count()
 
         f.error = Error("error")
@@ -904,8 +902,7 @@ class TestFlowMaster:
         assert not fm.do_server_playback(r)
         assert fm.do_server_playback(tutils.tflow())
 
-        q = queue.Queue()
-        fm.tick(q, 0)
+        fm.tick(0)
         assert fm.should_exit.is_set()
 
         fm.stop_server_playback()
diff --git a/test/mitmproxy/test_proxy.py b/test/mitmproxy/test_proxy.py
index fddb851e..e0897135 100644
--- a/test/mitmproxy/test_proxy.py
+++ b/test/mitmproxy/test_proxy.py
@@ -175,7 +175,7 @@ class TestDummyServer:
 
     def test_simple(self):
         d = DummyServer(None)
-        d.start_slave()
+        d.set_channel(None)
         d.shutdown()
 
 
diff --git a/test/mitmproxy/test_server.py b/test/mitmproxy/test_server.py
index 0d56e7ea..8843ee62 100644
--- a/test/mitmproxy/test_server.py
+++ b/test/mitmproxy/test_server.py
@@ -13,7 +13,7 @@ from netlib.tutils import raises
 from pathod import pathoc, pathod
 
 from mitmproxy.proxy.config import HostMatcher
-from mitmproxy.protocol import Kill
+from mitmproxy.exceptions import Kill
 from mitmproxy.models import Error, HTTPResponse
 
 from . import tutils, tservers
@@ -126,7 +126,7 @@ class TcpMixin:
         i2 = self.pathod("306")
         self._ignore_off()
 
-        self.master.masterq.join()
+        self.master.event_queue.join()
 
         assert n.status_code == 304
         assert i.status_code == 305
@@ -172,7 +172,7 @@ class TcpMixin:
         i2 = self.pathod("306")
         self._tcpproxy_off()
 
-        self.master.masterq.join()
+        self.master.event_queue.join()
 
         assert n.status_code == 304
         assert i.status_code == 305