3 files changed, 42 insertions, 1 deletions
diff --git a/test/mitmproxy/addons/test_save.py b/test/mitmproxy/addons/test_save.py
index a4e425cd..2dee708f 100644
--- a/test/mitmproxy/addons/test_save.py
+++ b/test/mitmproxy/addons/test_save.py
@@ -44,6 +44,19 @@ def test_tcp(tmpdir):
         assert rd(p)
 
 
+def test_websocket(tmpdir):
+    sa = save.Save()
+    with taddons.context() as tctx:
+        p = str(tmpdir.join("foo"))
+        tctx.configure(sa, save_stream_file=p)
+
+        f = tflow.twebsocketflow()
+        sa.websocket_start(f)
+        sa.websocket_end(f)
+        tctx.configure(sa, save_stream_file=None)
+        assert rd(p)
+
+
 def test_save_command(tmpdir):
     sa = save.Save()
     with taddons.context() as tctx:
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)
diff --git a/test/mitmproxy/test_flowfilter.py b/test/mitmproxy/test_flowfilter.py
index c411258a..4eb37d81 100644
--- a/test/mitmproxy/test_flowfilter.py
+++ b/test/mitmproxy/test_flowfilter.py
@@ -420,6 +420,20 @@ class TestMatchingWebSocketFlow:
         e = self.err()
         assert self.q("~e", e)
 
+    def test_domain(self):
+        q = self.flow()
+        assert self.q("~d example.com", q)
+        assert not self.q("~d none", q)
+
+    def test_url(self):
+        q = self.flow()
+        assert self.q("~u example.com", q)
+        assert self.q("~u example.com/ws", q)
+        assert not self.q("~u moo/path", q)
+
+        q.handshake_flow = None
+        assert not self.q("~u example.com", q)
+
     def test_body(self):
         f = self.flow()