websocket: test killable

1 files changed, 15 insertions, 0 deletions

diff --git a/test/mitmproxy/test_websocket.py b/test/mitmproxy/test_websocket.py
index 0a37f370..fcacec36 100644
--- a/test/mitmproxy/test_websocket.py
+++ b/test/mitmproxy/test_websocket.py
@@ -3,6 +3,7 @@ import pytest
 
 from mitmproxy.io import tnetstring
 from mitmproxy import flowfilter
+from mitmproxy.exceptions import Kill, ControlException
 from mitmproxy.test import tflow
 
 
@@ -42,6 +43,20 @@ class TestWebSocketFlow:
         assert f.error.get_state() == f2.error.get_state()
         assert f.error is not f2.error
 
+    def test_kill(self):
+        f = tflow.twebsocketflow()
+        with pytest.raises(ControlException):
+            f.intercept()
+            f.resume()
+            f.kill()
+
+        f = tflow.twebsocketflow()
+        f.intercept()
+        assert f.killable
+        f.kill()
+        assert not f.killable
+        assert f.reply.value == Kill
+
     def test_match(self):
         f = tflow.twebsocketflow()
         assert not flowfilter.match("~b nonexistent", f)