From 59c277effd016fdf771a65aaaa87526ff5fe869e Mon Sep 17 00:00:00 2001 From: Thomas Kriechbaumer Date: Fri, 29 Dec 2017 21:55:41 +0100 Subject: websocket: add flow kill test --- test/mitmproxy/proxy/protocol/test_websocket.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) (limited to 'test') diff --git a/test/mitmproxy/proxy/protocol/test_websocket.py b/test/mitmproxy/proxy/protocol/test_websocket.py index a7acdc4d..490d31a0 100644 --- a/test/mitmproxy/proxy/protocol/test_websocket.py +++ b/test/mitmproxy/proxy/protocol/test_websocket.py @@ -221,6 +221,25 @@ class TestSimple(_WebSocketTest): assert frame.payload == b'foo' +class TestKillFlow(_WebSocketTest): + + @classmethod + def handle_websockets(cls, rfile, wfile): + wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.TEXT, payload=b'server-foobar'))) + wfile.flush() + + def test_kill(self): + class KillFlow: + def websocket_message(self, f): + f.kill() + + self.master.addons.add(KillFlow()) + self.setup_connection() + + frame = websockets.Frame.from_file(self.client.rfile) + assert frame.payload == b'foo' + + class TestSimpleTLS(_WebSocketTest): ssl = True -- cgit v1.2.3