Merge pull request #1955 from Kriechi/test++

increase addon test coverage

2 files changed, 36 insertions, 1 deletions

diff --git a/test/mitmproxy/addons/test_disable_h2c_upgrade.py b/test/mitmproxy/addons/test_disable_h2c_upgrade.py
new file mode 100644
index 00000000..6cab713d
--- /dev/null
+++ b/test/mitmproxy/addons/test_disable_h2c_upgrade.py
@@ -0,0 +1,17 @@
+from mitmproxy.addons import disable_h2c_upgrade
+from mitmproxy.test import tflow
+
+
+class TestTermLog:
+    def test_simple(self):
+        a = disable_h2c_upgrade.DisableH2CleartextUpgrade()
+
+        f = tflow.tflow()
+        f.request.headers['upgrade'] = 'h2c'
+        f.request.headers['connection'] = 'foo'
+        f.request.headers['http2-settings'] = 'bar'
+
+        a.request(f)
+        assert 'upgrade' not in f.request.headers
+        assert 'connection' not in f.request.headers
+        assert 'http2-settings' not in f.request.headers
diff --git a/test/mitmproxy/addons/test_dumper.py b/test/mitmproxy/addons/test_dumper.py
index 7ac17a7c..8fa8a22a 100644
--- a/test/mitmproxy/addons/test_dumper.py
+++ b/test/mitmproxy/addons/test_dumper.py
@@ -157,7 +157,7 @@ def test_tcp():
     d = dumper.Dumper(sio)
     with taddons.context(options=dump.Options()) as ctx:
         ctx.configure(d, flow_detail=3, showhost=True)
-        f = tflow.ttcpflow(client_conn=True, server_conn=True)
+        f = tflow.ttcpflow()
         d.tcp_message(f)
         assert "it's me" in sio.getvalue()
         sio.truncate(0)
@@ -165,3 +165,21 @@ def test_tcp():
         f = tflow.ttcpflow(client_conn=True, err=True)
         d.tcp_error(f)
         assert "Error in TCP" in sio.getvalue()
+
+
+def test_websocket():
+    sio = io.StringIO()
+    d = dumper.Dumper(sio)
+    with taddons.context(options=dump.Options()) as ctx:
+        ctx.configure(d, flow_detail=3, showhost=True)
+        f = tflow.twebsocketflow()
+        d.websocket_message(f)
+        assert "hello text" in sio.getvalue()
+        sio.truncate(0)
+
+        d.websocket_end(f)
+        assert "WebSocket connection closed by" in sio.getvalue()
+
+        f = tflow.twebsocketflow(client_conn=True, err=True)
+        d.websocket_error(f)
+        assert "Error in WebSocket" in sio.getvalue()