add tests

5 files changed, 33 insertions, 3 deletions

diff --git a/mitmproxy/contrib/wsproto/__init__.py b/mitmproxy/contrib/wsproto/__init__.py
index e69de29b..d0592bc5 100644
--- a/mitmproxy/contrib/wsproto/__init__.py
+++ b/mitmproxy/contrib/wsproto/__init__.py
@@ -0,0 +1,13 @@
+from . import compat
+from . import connection
+from . import events
+from . import extensions
+from . import frame_protocol
+
+__all__ = [
+    'compat',
+    'connection',
+    'events',
+    'extensions',
+    'frame_protocol',
+]
diff --git a/mitmproxy/test/tflow.py b/mitmproxy/test/tflow.py
index cb2b940a..91747866 100644
--- a/mitmproxy/test/tflow.py
+++ b/mitmproxy/test/tflow.py
@@ -75,7 +75,9 @@ def twebsocketflow(client_conn=True, server_conn=True, messages=True, err=None,
         handshake_flow.response = resp
 
     f = websocket.WebSocketFlow(client_conn, server_conn, handshake_flow)
-    handshake_flow.metadata['websocket_flow'] = f
+    f.metadata['websocket_handshake'] = handshake_flow.id
+    handshake_flow.metadata['websocket_flow'] = f.id
+    handshake_flow.metadata['websocket'] = True
 
     if messages is True:
         messages = [
diff --git a/mitmproxy/websocket.py b/mitmproxy/websocket.py
index ade23732..8efd4117 100644
--- a/mitmproxy/websocket.py
+++ b/mitmproxy/websocket.py
@@ -13,7 +13,7 @@ class WebSocketMessage(serializable.Serializable):
     def __init__(
         self, type: int, from_client: bool, content: bytes, timestamp: Optional[int]=None
     ) -> None:
-        self.type = wsproto.frame_protocol.Opcode(type)
+        self.type = wsproto.frame_protocol.Opcode(type)  # type: ignore
         self.from_client = from_client
         self.content = content
         self.timestamp = timestamp or int(time.time())  # type: int
diff --git a/test/mitmproxy/addons/test_save.py b/test/mitmproxy/addons/test_save.py
index 84564157..2dee708f 100644
--- a/test/mitmproxy/addons/test_save.py
+++ b/test/mitmproxy/addons/test_save.py
@@ -52,6 +52,7 @@ def test_websocket(tmpdir):
 
         f = tflow.twebsocketflow()
         sa.websocket_start(f)
+        sa.websocket_end(f)
         tctx.configure(sa, save_stream_file=None)
         assert rd(p)
 
diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py
index fcc766b5..8cc11a16 100644
--- a/test/mitmproxy/test_flow.py
+++ b/test/mitmproxy/test_flow.py
@@ -97,7 +97,7 @@ class TestSerialize:
 
 
 class TestFlowMaster:
-    def test_load_flow_reverse(self):
+    def test_load_http_flow_reverse(self):
         s = tservers.TestState()
         opts = options.Options(
             mode="reverse:https://use-this-domain"
@@ -108,6 +108,20 @@ class TestFlowMaster:
         fm.load_flow(f)
         assert s.flows[0].request.host == "use-this-domain"
 
+    def test_load_websocket_flow(self):
+        s = tservers.TestState()
+        opts = options.Options(
+            mode="reverse:https://use-this-domain"
+        )
+        fm = master.Master(opts)
+        fm.addons.add(s)
+        f = tflow.twebsocketflow()
+        fm.load_flow(f.handshake_flow)
+        fm.load_flow(f)
+        assert s.flows[0].request.host == "use-this-domain"
+        assert s.flows[1].handshake_flow == f.handshake_flow
+        assert len(s.flows[1].messages) == len(f.messages)
+
     def test_replay(self):
         opts = options.Options()
         fm = master.Master(opts)