diff options
Diffstat (limited to 'test')
| -rw-r--r-- | test/mitmproxy/protocol/test_websocket.py | 28 | ||||
| -rw-r--r-- | test/mitmproxy/tservers.py | 8 |
2 files changed, 36 insertions, 0 deletions
diff --git a/test/mitmproxy/protocol/test_websocket.py b/test/mitmproxy/protocol/test_websocket.py index c5db9341..e1c3e49a 100644 --- a/test/mitmproxy/protocol/test_websocket.py +++ b/test/mitmproxy/protocol/test_websocket.py @@ -5,6 +5,8 @@ import traceback from mitmproxy import options from mitmproxy import exceptions +from mitmproxy.http import HTTPFlow +from mitmproxy.websocket import WebSocketFlow from mitmproxy.proxy.config import ProxyConfig import mitmproxy.net @@ -147,6 +149,10 @@ class TestSimple(_WebSocketTest): wfile.write(bytes(frame)) wfile.flush() + frame = websockets.Frame.from_file(rfile) + wfile.write(bytes(frame)) + wfile.flush() + def test_simple(self): client = self._setup_connection() @@ -159,9 +165,31 @@ class TestSimple(_WebSocketTest): frame = websockets.Frame.from_file(client.rfile) assert frame.payload == b'client-foobar' + client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.BINARY, payload=b'\xde\xad\xbe\xef'))) + client.wfile.flush() + + frame = websockets.Frame.from_file(client.rfile) + assert frame.payload == b'\xde\xad\xbe\xef' + client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.CLOSE))) client.wfile.flush() + assert len(self.master.state.flows) == 2 + assert isinstance(self.master.state.flows[0], HTTPFlow) + 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[1].content == b'client-foobar' + assert self.master.state.flows[1].messages[1].type == '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[3].content == b'\xde\xad\xbe\xef' + assert self.master.state.flows[1].messages[3].type == '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] + class TestSimpleTLS(_WebSocketTest): ssl = True diff --git a/test/mitmproxy/tservers.py b/test/mitmproxy/tservers.py index f9dfde30..1020f23c 100644 --- a/test/mitmproxy/tservers.py +++ b/test/mitmproxy/tservers.py @@ -26,6 +26,14 @@ class TestState: if f not in self.flows: self.flows.append(f) + def websocket_start(self, f): + if f not in self.flows: + self.flows.append(f) + + def tcp_start(self, f): + if f not in self.flows: + self.flows.append(f) + # FIXME: compat with old state - remove in favor of len(state.flows) def flow_count(self): return len(self.flows) |