From be84049cc2413a76754730fad3fe70ca72ecab03 Mon Sep 17 00:00:00 2001 From: Greg Femec Date: Fri, 28 Jul 2017 21:28:02 -0700 Subject: [PATCH 1/2] add message parsing test --- test/test_websocket.py | 54 +++++++++++++++++++++++++++++++++++++++++- 1 file changed, 53 insertions(+), 1 deletion(-) diff --git a/test/test_websocket.py b/test/test_websocket.py index 96fb384..d88518f 100644 --- a/test/test_websocket.py +++ b/test/test_websocket.py @@ -1,6 +1,7 @@ # -*- coding: utf-8 -*- import unittest import os +import ssl import socket import struct @@ -135,7 +136,58 @@ def test_send_message_without_masking(self): ws = WebSocket(sock=m) ws.send(tm) m.sendall.assert_called_once_with(tm.single()) - + + def message_parsing_test(self, messages, use_ssl=False): + to_send = [ + b"".join(TextMessage(msg).single(mask=True) for msg in messages) + ] + + def recv(max_bytes): + data = to_send[0][:max_bytes] + to_send[0] = to_send[0][max_bytes:] + return data + + recv_mock = MagicMock(side_effect=recv) + + if use_ssl: + pending_mock = MagicMock(side_effect=lambda: len(to_send[0])) + socket_mock = MagicMock(spec=ssl.SSLSocket) + socket_patch = patch.multiple(socket_mock, recv=recv_mock, + pending=pending_mock) + else: + socket_mock = MagicMock(spec=socket.socket) + socket_patch = patch.multiple(socket_mock, recv=recv_mock) + + index = [0] + + def received_message(msg): + self.assertEqual(msg.data, messages[index[0]]) + index[0] += 1 + + received_message_mock = MagicMock(side_effect=received_message) + + websocket = WebSocket(sock=socket_mock) + websocket._is_secure = use_ssl + websocket_patch = \ + patch.multiple(websocket, received_message=received_message_mock) + with socket_patch, websocket_patch: + while websocket.once(): + pass + + self.assertEqual(index[0], len(messages)) + + def test_message_parsing(self): + self.message_parsing_test([b'hello'], False) + + def test_message_parsing_ssl(self): + self.message_parsing_test([b'hello'], True) + + def test_messages_parsing(self): + self.message_parsing_test([b'hello', b'world'], False) + + def test_messages_parsing_ssl(self): + self.message_parsing_test([b'hello', b'world'], True) + def test_send_generator_without_masking(self): tm0 = b'hello' tm1 = b'world' From 2121b34ced26d75e103e80df14ab7429349851e2 Mon Sep 17 00:00:00 2001 From: Greg Femec Date: Thu, 27 Jul 2017 21:54:59 -0700 Subject: [PATCH 2/2] process buffer in reading_buffer_size chunks --- ws4py/websocket.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/ws4py/websocket.py b/ws4py/websocket.py index 511921e..085105a 100644 --- a/ws4py/websocket.py +++ b/ws4py/websocket.py @@ -400,12 +400,12 @@ def once(self): self.unhandled_error(e) return False else: - # process as much as we can - # the process will stop either if there is no buffer left - # or if the stream is closed - if not self.process(self.buf): - return False - self.buf = b"" + # process buffer in reading_buffer_size chunks + while len(self.buf) > 0 and self.reading_buffer_size > 0: + b = self.buf[:self.reading_buffer_size] + self.buf = self.buf[self.reading_buffer_size:] + if not self.process(b): + return False return True