|
4 | 4 | import socket
|
5 | 5 | import struct
|
6 | 6 |
|
| 7 | +try: |
| 8 | + from io import BytesIO |
| 9 | +except ImportError: |
| 10 | + from StringIO import StringIO as BytesIO |
| 11 | + |
7 | 12 | from mock import MagicMock, call, patch
|
8 | 13 |
|
9 | 14 | from ws4py.framing import Frame, \
|
@@ -178,24 +183,23 @@ def test_sending_ping(self):
|
178 | 183 | m.sendall.assert_called_once_with(tm)
|
179 | 184 |
|
180 | 185 | 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) |
199 | 203 |
|
200 | 204 |
|
201 | 205 | if __name__ == '__main__':
|
|
0 commit comments