Skip to content

Commit 0714fb2

Browse files
committed
Make test_websocket.test_spill_frame more readable
1 parent aa644bf commit 0714fb2

File tree

1 file changed

+22
-18
lines changed

1 file changed

+22
-18
lines changed

test/test_websocket.py

Lines changed: 22 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,11 @@
44
import socket
55
import struct
66

7+
try:
8+
from io import BytesIO
9+
except ImportError:
10+
from StringIO import StringIO as BytesIO
11+
712
from mock import MagicMock, call, patch
813

914
from ws4py.framing import Frame, \
@@ -178,24 +183,23 @@ def test_sending_ping(self):
178183
m.sendall.assert_called_once_with(tm)
179184

180185
def test_spill_frame(self):
181-
s = MagicMock()
182-
m = MagicMock()
183-
c = MagicMock()
184-
recv = lambda size: b'a' * size
185-
186-
with patch.multiple(m, recv=recv):
187-
ws = WebSocket(sock=m)
188-
sz = 20
189-
spill = 10
190-
proc = MagicMock(return_value=('a' * sz))
191-
pend = lambda: b'a' * spill
192-
193-
with patch.multiple(ws, close=c, process=proc, _get_from_pending=pend):
194-
ws.stream = s
195-
ws.reading_buffer_size = sz
196-
ws.once()
197-
proc.assert_called_once_with(b'a' * sz)
198-
self.assertEqual(len(ws.buf), spill)
186+
data = b"hello"
187+
buf = BytesIO(data + b"spillover")
188+
189+
sock = MagicMock()
190+
sock._ssl = object() # for WebSocket._is_secure logic
191+
sock.recv.side_effect = buf.read
192+
sock.pending.side_effect = lambda: buf.tell() < len(buf.getvalue())
193+
194+
ws = WebSocket(sock=sock)
195+
ws.stream = MagicMock()
196+
197+
self.assertTrue(ws._is_secure)
198+
199+
ws.reading_buffer_size = len(data)
200+
ws.once()
201+
202+
ws.stream.parser.send.assert_called_once_with(data)
199203

200204

201205
if __name__ == '__main__':

0 commit comments

Comments
 (0)