From a2d45193546962e2e14d1959e1bf008c83b9f3cf Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Sat, 24 Mar 2018 12:03:50 +1300 Subject: asyncio: brutally rip out our old queue mechanism --- test/mitmproxy/proxy/protocol/test_http2.py | 4 +- test/mitmproxy/proxy/protocol/test_websocket.py | 4 +- test/mitmproxy/proxy/test_server.py | 8 -- test/mitmproxy/test_controller.py | 140 ++++++++++++++---------- test/mitmproxy/test_flow.py | 7 +- test/mitmproxy/tservers.py | 40 +++---- 6 files changed, 106 insertions(+), 97 deletions(-) (limited to 'test') 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) -- cgit v1.2.3