Merge pull request #2732 from mitmproxy/websocket-flow-kill

fix flow.kill

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..d9389faf 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()
+
+        with pytest.raises(exceptions.TcpDisconnect):
+            websockets.Frame.from_file(self.client.rfile)
+
+
 class TestSimpleTLS(_WebSocketTest):
     ssl = True