websockets, tcp, version: coverage++

3 files changed, 102 insertions, 7 deletions

diff --git a/test/mitmproxy/addons/test_stickycookie.py b/test/mitmproxy/addons/test_stickycookie.py
index 157f2959..9092e09b 100644
--- a/test/mitmproxy/addons/test_stickycookie.py
+++ b/test/mitmproxy/addons/test_stickycookie.py
@@ -39,7 +39,6 @@ class TestStickyCookie:
             assert "cookie" not in f.request.headers
 
             f = f.copy()
-            f.reply.acked = False
             sc.request(f)
             assert f.request.headers["cookie"] == "foo=bar"
 
diff --git a/test/mitmproxy/protocol/test_websocket.py b/test/mitmproxy/protocol/test_websocket.py
index e42250e0..73ee8b35 100644
--- a/test/mitmproxy/protocol/test_websocket.py
+++ b/test/mitmproxy/protocol/test_websocket.py
@@ -179,16 +179,15 @@ class TestSimple(_WebSocketTest):
         assert isinstance(self.master.state.flows[1], WebSocketFlow)
         assert len(self.master.state.flows[1].messages) == 5
         assert self.master.state.flows[1].messages[0].content == b'server-foobar'
-        assert self.master.state.flows[1].messages[0].type == 'text'
+        assert self.master.state.flows[1].messages[0].type == websockets.OPCODE.TEXT
         assert self.master.state.flows[1].messages[1].content == b'client-foobar'
-        assert self.master.state.flows[1].messages[1].type == 'text'
+        assert self.master.state.flows[1].messages[1].type == websockets.OPCODE.TEXT
         assert self.master.state.flows[1].messages[2].content == b'client-foobar'
-        assert self.master.state.flows[1].messages[2].type == 'text'
+        assert self.master.state.flows[1].messages[2].type == websockets.OPCODE.TEXT
         assert self.master.state.flows[1].messages[3].content == b'\xde\xad\xbe\xef'
-        assert self.master.state.flows[1].messages[3].type == 'binary'
+        assert self.master.state.flows[1].messages[3].type == websockets.OPCODE.BINARY
         assert self.master.state.flows[1].messages[4].content == b'\xde\xad\xbe\xef'
-        assert self.master.state.flows[1].messages[4].type == 'binary'
-        assert [m.info for m in self.master.state.flows[1].messages]
+        assert self.master.state.flows[1].messages[4].type == websockets.OPCODE.BINARY
 
 
 class TestSimpleTLS(_WebSocketTest):
diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py
index 371474ff..65e6845f 100644
--- a/test/mitmproxy/test_flow.py
+++ b/test/mitmproxy/test_flow.py
@@ -10,6 +10,7 @@ from mitmproxy.exceptions import FlowReadException, Kill
 from mitmproxy import flow
 from mitmproxy import http
 from mitmproxy import connections
+from mitmproxy import tcp
 from mitmproxy.proxy import ProxyConfig
 from mitmproxy.proxy.server import DummyServer
 from mitmproxy import master
@@ -156,8 +157,99 @@ class TestHTTPFlow:
         assert f.response.raw_content == b"abarb"
 
 
+class TestWebSocketFlow:
+
+    def test_copy(self):
+        f = tflow.twebsocketflow()
+        f.get_state()
+        f2 = f.copy()
+        a = f.get_state()
+        b = f2.get_state()
+        del a["id"]
+        del b["id"]
+        del a["handshake_flow"]["id"]
+        del b["handshake_flow"]["id"]
+        assert a == b
+        assert not f == f2
+        assert f is not f2
+
+        assert f.client_key == f2.client_key
+        assert f.client_protocol == f2.client_protocol
+        assert f.client_extensions == f2.client_extensions
+        assert f.server_accept == f2.server_accept
+        assert f.server_protocol == f2.server_protocol
+        assert f.server_extensions == f2.server_extensions
+        assert f.messages is not f2.messages
+        assert f.handshake_flow is not f2.handshake_flow
+
+        for m in f.messages:
+            m2 = m.copy()
+            m2.set_state(m2.get_state())
+            assert m is not m2
+            assert m.get_state() == m2.get_state()
+
+        f = tflow.twebsocketflow(err=True)
+        f2 = f.copy()
+        assert f is not f2
+        assert f.handshake_flow is not f2.handshake_flow
+        assert f.error.get_state() == f2.error.get_state()
+        assert f.error is not f2.error
+
+    def test_match(self):
+        f = tflow.twebsocketflow()
+        assert not flowfilter.match("~b nonexistent", f)
+        assert flowfilter.match(None, f)
+        assert not flowfilter.match("~b nonexistent", f)
+
+        f = tflow.twebsocketflow(err=True)
+        assert flowfilter.match("~e", f)
+
+        with pytest.raises(ValueError):
+            flowfilter.match("~", f)
+
+    def test_repr(self):
+        f = tflow.twebsocketflow()
+        assert 'WebSocketFlow' in repr(f)
+        assert 'binary message: ' in repr(f.messages[0])
+        assert 'text message: ' in repr(f.messages[1])
+
+
 class TestTCPFlow:
 
+    def test_copy(self):
+        f = tflow.ttcpflow()
+        f.get_state()
+        f2 = f.copy()
+        a = f.get_state()
+        b = f2.get_state()
+        del a["id"]
+        del b["id"]
+        assert a == b
+        assert not f == f2
+        assert f is not f2
+
+        assert f.messages is not f2.messages
+
+        for m in f.messages:
+            assert m.get_state()
+            m2 = m.copy()
+            assert not m == m2
+            assert m is not m2
+
+            a = m.get_state()
+            b = m2.get_state()
+            assert a == b
+
+        m = tcp.TCPMessage(False, 'foo')
+        m.set_state(f.messages[0].get_state())
+        assert m.timestamp == f.messages[0].timestamp
+
+        f = tflow.ttcpflow(err=True)
+        f2 = f.copy()
+        assert f is not f2
+        assert f.error.get_state() == f2.error.get_state()
+        assert f.error is not f2.error
+
     def test_match(self):
         f = tflow.ttcpflow()
         assert not flowfilter.match("~b nonexistent", f)
@@ -170,6 +262,11 @@ class TestTCPFlow:
         with pytest.raises(ValueError):
             flowfilter.match("~", f)
 
+    def test_repr(self):
+        f = tflow.ttcpflow()
+        assert 'TCPFlow' in repr(f)
+        assert '-> ' in repr(f.messages[0])
+
 
 class TestSerialize: