From d674de298cbaef2eb5348ed65ef51775edadb89b Mon Sep 17 00:00:00 2001
From: Thomas Kriechbaumer <thomas@kriechbaumer.name>
Date: Mon, 23 Jan 2017 15:45:53 +0100
Subject: [PATCH] test websocket addon

---
 mitmproxy/test/tflow.py              | 55 ++++++++++++++++++++++++++++
 test/mitmproxy/addons/test_dumper.py | 19 +++++++++-
 2 files changed, 73 insertions(+), 1 deletion(-)

diff --git a/mitmproxy/test/tflow.py b/mitmproxy/test/tflow.py
index 959c9a2c3..a93c86ddd 100644
--- a/mitmproxy/test/tflow.py
+++ b/mitmproxy/test/tflow.py
@@ -1,9 +1,11 @@
 from mitmproxy.test import tutils
 from mitmproxy import tcp
+from mitmproxy import websocket
 from mitmproxy import controller
 from mitmproxy import http
 from mitmproxy import connections
 from mitmproxy import flow
+from mitmproxy.net import http as net_http
 
 
 def ttcpflow(client_conn=True, server_conn=True, messages=True, err=None):
@@ -25,6 +27,59 @@ def ttcpflow(client_conn=True, server_conn=True, messages=True, err=None):
     f.reply = controller.DummyReply()
     return f
 
+def twebsocketflow(client_conn=True, server_conn=True, messages=True, err=None, handshake_flow=True):
+
+    if client_conn is True:
+        client_conn = tclient_conn()
+    if server_conn is True:
+        server_conn = tserver_conn()
+    if handshake_flow is True:
+        req = http.HTTPRequest(
+            "relative",
+            "GET",
+            "http",
+            "example.com",
+            "80",
+            "/ws",
+            "HTTP/1.1",
+            headers=net_http.Headers(
+                connection="upgrade",
+                upgrade="websocket",
+                sec_websocket_version="13",
+                sec_websocket_key="1234",
+            ),
+            content=b''
+        )
+        resp = http.HTTPResponse(
+            "HTTP/1.1",
+            101,
+            reason=net_http.status_codes.RESPONSES.get(101),
+            headers=net_http.Headers(
+                connection='upgrade',
+                upgrade='websocket',
+                sec_websocket_accept=b'',
+            ),
+            content=b'',
+        )
+        handshake_flow = http.HTTPFlow(client_conn, server_conn)
+        handshake_flow.request = req
+        handshake_flow.response = resp
+
+    f = websocket.WebSocketFlow(client_conn, server_conn, handshake_flow)
+
+    if messages is True:
+        messages = [
+            websocket.WebSocketBinaryMessage(f, True, b"hello binary"),
+            websocket.WebSocketTextMessage(f, False, "hello text".encode()),
+        ]
+    if err is True:
+        err = terr()
+
+    f.messages = messages
+    f.error = err
+    f.reply = controller.DummyReply()
+    return f
+
 
 def tflow(client_conn=True, server_conn=True, req=True, resp=None, err=None):
     """
diff --git a/test/mitmproxy/addons/test_dumper.py b/test/mitmproxy/addons/test_dumper.py
index 7ac17a7c5..1a2639a2d 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,20 @@ 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()