3 files changed, 42 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
 
diff --git a/test/mitmproxy/test_websocket.py b/test/mitmproxy/test_websocket.py
index 7c53a4b0..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)
@@ -71,3 +86,9 @@ class TestWebSocketFlow:
         d = tflow.twebsocketflow().handshake_flow.get_state()
         tnetstring.dump(d, b)
         assert b.getvalue()
+
+    def test_message_kill(self):
+        f = tflow.twebsocketflow()
+        assert not f.messages[-1].killed
+        f.messages[-1].kill()
+        assert f.messages[-1].killed
diff --git a/test/mitmproxy/utils/test_human.py b/test/mitmproxy/utils/test_human.py
index e8ffaad4..947cfa4a 100644
--- a/test/mitmproxy/utils/test_human.py
+++ b/test/mitmproxy/utils/test_human.py
@@ -54,3 +54,5 @@ def test_format_address():
     assert human.format_address(("::ffff:127.0.0.1", "54010", "0", "0")) == "127.0.0.1:54010"
     assert human.format_address(("127.0.0.1", "54010")) == "127.0.0.1:54010"
     assert human.format_address(("example.com", "54010")) == "example.com:54010"
+    assert human.format_address(("::", "8080")) == "*:8080"
+    assert human.format_address(("0.0.0.0", "8080")) == "*:8080"