4 files changed, 65 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/tools/console/test_defaultkeys.py b/test/mitmproxy/tools/console/test_defaultkeys.py
new file mode 100644
index 00000000..1f17c888
--- /dev/null
+++ b/test/mitmproxy/tools/console/test_defaultkeys.py
@@ -0,0 +1,23 @@
+from mitmproxy.test.tflow import tflow
+from mitmproxy.tools.console import defaultkeys
+from mitmproxy.tools.console import keymap
+from mitmproxy.tools.console import master
+from mitmproxy import command
+
+
+def test_commands_exist():
+    km = keymap.Keymap(None)
+    defaultkeys.map(km)
+    assert km.bindings
+    m = master.ConsoleMaster(None)
+    m.load_flow(tflow())
+
+    for binding in km.bindings:
+        cmd, *args = command.lexer(binding.command)
+        assert cmd in m.commands.commands
+
+        cmd_obj = m.commands.commands[cmd]
+        try:
+            cmd_obj.prepare_args(args)
+        except Exception as e:
+            raise ValueError("Invalid command: {}".format(binding.command)) from e
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"