From 793f5c5681ade96d6650130361229dd05b6a05b0 Mon Sep 17 00:00:00 2001
From: Alexander Prinzhorn <alexander@prinzhorn.it>
Date: Tue, 13 Jul 2021 12:51:24 +0200
Subject: [PATCH] fix dumper test cov

---
 mitmproxy/test/tflow.py              | 18 ++++++++++++------
 test/mitmproxy/addons/test_dumper.py | 20 ++++++++++++++++++++
 2 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/mitmproxy/test/tflow.py b/mitmproxy/test/tflow.py
index 10d6e03cd..615e92434 100644
--- a/mitmproxy/test/tflow.py
+++ b/mitmproxy/test/tflow.py
@@ -30,7 +30,7 @@ def ttcpflow(client_conn=True, server_conn=True, messages=True, err=None):
     return f
 
 
-def twebsocketflow(messages=True, err=None) -> http.HTTPFlow:
+def twebsocketflow(messages=True, err=None, close_code=None, close_reason='') -> http.HTTPFlow:
     flow = http.HTTPFlow(tclient_conn(), tserver_conn())
     flow.request = http.Request(
         "example.com",
@@ -74,12 +74,18 @@ def twebsocketflow(messages=True, err=None) -> http.HTTPFlow:
             websocket.WebSocketMessage(Opcode.TEXT, True, b"hello text", 946681204),
             websocket.WebSocketMessage(Opcode.TEXT, False, b"it's me", 946681205),
         ]
-    if err is True:
-        # ABNORMAL_CLOSURE
-        flow.websocket.close_code = 1006
+
+    flow.websocket.close_reason = close_reason
+
+    if close_code is not None:
+        flow.websocket.close_code = close_code
     else:
-        # NORMAL_CLOSURE
-        flow.websocket.close_code = 1000
+        if err is True:
+            # ABNORMAL_CLOSURE
+            flow.websocket.close_code = 1006
+        else:
+            # NORMAL_CLOSURE
+            flow.websocket.close_code = 1000
 
     flow.reply = controller.DummyReply()
     return flow
diff --git a/test/mitmproxy/addons/test_dumper.py b/test/mitmproxy/addons/test_dumper.py
index bffe721f2..2f6a86746 100644
--- a/test/mitmproxy/addons/test_dumper.py
+++ b/test/mitmproxy/addons/test_dumper.py
@@ -232,10 +232,30 @@ def test_websocket():
 
         d.websocket_end(f)
         assert "WebSocket connection closed by" in sio.getvalue()
+        sio_err.truncate(0)
 
         f = tflow.twebsocketflow(err=True)
         d.websocket_end(f)
         assert "Error in WebSocket" in sio_err.getvalue()
+        assert "(reason:" not in sio_err.getvalue()
+        sio_err.truncate(0)
+
+        f = tflow.twebsocketflow(err=True, close_reason='Some lame excuse')
+        d.websocket_end(f)
+        assert "Error in WebSocket" in sio_err.getvalue()
+        assert "(reason: Some lame excuse)" in sio_err.getvalue()
+        sio_err.truncate(0)
+
+        f = tflow.twebsocketflow(close_code=4000)
+        d.websocket_end(f)
+        assert "UNKNOWN_ERROR=4000" in sio_err.getvalue()
+        assert "(reason:" not in sio_err.getvalue()
+        sio_err.truncate(0)
+
+        f = tflow.twebsocketflow(close_code=4000, close_reason='I swear I had a reason')
+        d.websocket_end(f)
+        assert "UNKNOWN_ERROR=4000" in sio_err.getvalue()
+        assert "(reason: I swear I had a reason)" in sio_err.getvalue()
 
 
 def test_http2():