-
-
Notifications
You must be signed in to change notification settings - Fork 3.9k
/
test_websocket.py
63 lines (50 loc) · 2.06 KB
/
test_websocket.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
import pytest
from wsproto.frame_protocol import Opcode
from mitmproxy import http
from mitmproxy import websocket
from mitmproxy.test import tflow
class TestWebSocketData:
def test_repr(self):
assert repr(tflow.twebsocketflow().websocket) == "<WebSocketData (3 messages)>"
def test_state(self):
f = tflow.twebsocketflow()
f2 = http.HTTPFlow.from_state(f.get_state())
f2.set_state(f.get_state())
def test_formatting(self):
tf = tflow.twebsocketflow().websocket
formatted_messages = tf._get_formatted_messages()
assert b"[OUTGOING] hello binary" in formatted_messages
assert b"[OUTGOING] hello text" in formatted_messages
assert b"[INCOMING] it's me" in formatted_messages
class TestWebSocketMessage:
def test_basic(self):
m = websocket.WebSocketMessage(Opcode.TEXT, True, b"foo")
m.set_state(m.get_state())
assert m.content == b"foo"
assert repr(m) == "'foo'"
m.type = Opcode.BINARY
assert repr(m) == "b'foo'"
assert not m.dropped
m.drop()
assert m.dropped
def test_text(self):
txt = websocket.WebSocketMessage(Opcode.TEXT, True, b"foo")
bin = websocket.WebSocketMessage(Opcode.BINARY, True, b"foo")
assert txt.is_text
assert txt.text == "foo"
txt.text = "bar"
assert txt.content == b"bar"
assert not bin.is_text
with pytest.raises(AttributeError, match="do not have a 'text' attribute."):
_ = bin.text
with pytest.raises(AttributeError, match="do not have a 'text' attribute."):
bin.text = "bar"
def test_message_formatting(self):
incoming_message = websocket.WebSocketMessage(
Opcode.BINARY, False, b"Test Incoming"
)
outgoing_message = websocket.WebSocketMessage(
Opcode.BINARY, True, b"Test OutGoing"
)
assert incoming_message._format_ws_message() == b"[INCOMING] Test Incoming"
assert outgoing_message._format_ws_message() == b"[OUTGOING] Test OutGoing"