diff options
| -rw-r--r-- | test/mitmproxy/proxy/protocol/test_websocket.py | 19 |
1 files changed, 19 insertions, 0 deletions
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 |