diff --git a/test/test_websocket.py b/test/test_websocket.py index 9f22242..501218d 100644 --- a/test/test_websocket.py +++ b/test/test_websocket.py @@ -4,6 +4,11 @@ import socket import struct +try: + from io import BytesIO +except ImportError: + from StringIO import StringIO as BytesIO + from mock import MagicMock, call, patch from ws4py.framing import Frame, \ @@ -178,24 +183,23 @@ def test_sending_ping(self): m.sendall.assert_called_once_with(tm) def test_spill_frame(self): - s = MagicMock() - m = MagicMock() - c = MagicMock() - recv = lambda size: b'a' * size - - with patch.multiple(m, recv=recv): - ws = WebSocket(sock=m) - sz = 20 - spill = 10 - proc = MagicMock(return_value=('a' * sz)) - pend = lambda: b'a' * spill - - with patch.multiple(ws, close=c, process=proc, _get_from_pending=pend): - ws.stream = s - ws.reading_buffer_size = sz - ws.once() - proc.assert_called_once_with(b'a' * sz) - self.assertEqual(len(ws.buf), spill) + data = b"hello" + buf = BytesIO(data + b"spillover") + + sock = MagicMock() + sock._ssl = object() # for WebSocket._is_secure logic + sock.recv.side_effect = buf.read + sock.pending.side_effect = lambda: buf.tell() < len(buf.getvalue()) + + ws = WebSocket(sock=sock) + ws.stream = MagicMock() + + self.assertTrue(ws._is_secure) + + ws.reading_buffer_size = len(data) + ws.once() + + ws.stream.parser.send.assert_called_once_with(data) if __name__ == '__main__':