more unit test coverage of framing and websocket

Lawouach · Lawouach · commit fbb48de37f4b · 2013-10-20T00:07:27.000+02:00
diff --git a/test/test_frame.py b/test/test_frame.py
@@ -185,14 +185,93 @@ def test_fragmented_control_frame_is_invalid(self):
         f = Frame()
         self.assertRaises(ProtocolException, f.parser.send, b'0x9')
 
-    def test_frame_sized_below_127(self):
-        bytes = Frame(opcode=OPCODE_TEXT, body=b'*'*65536, fin=1).build()
+    def test_fragmented_control_frame_is_too_large(self):
+        bytes = Frame(opcode=OPCODE_PING, body=b'*'*65536, fin=1).build()
+        f = Frame()
+        self.assertRaises(FrameTooLargeException, f.parser.send, bytes)
+
+    def test_frame_sized_127(self):
+        body = b'*'*65536
+        bytes = Frame(opcode=OPCODE_TEXT, body=body, fin=1).build()
 
         f = Frame()
+        # determine how the size is stored
         f.parser.send(bytes[:3])
         self.assertTrue(f.masking_key is None)
+        # that's a large frame indeed
         self.assertEqual(f.payload_length, 127)
 
+        # this will compute the actual application data size
+        # it will also read the first byte of data
+        # indeed the length is found from byte 3 to 10
+        f.parser.send(bytes[3:11])
+        self.assertEqual(f.payload_length, 65536)
+        
+        # parse the rest of our data
+        f.parser.send(bytes[11:])
+        self.assertEqual(f.body, body)
+
+        
+        # The same but this time we provide enough
+        # bytes so that the application's data length
+        # can be computed from the first generator's send call
+        f = Frame()
+        f.parser.send(bytes[:10])
+        self.assertTrue(f.masking_key is None)
+        self.assertEqual(f.payload_length, 65536)
+        
+        # parse the rest of our data
+        f.parser.send(bytes[10:])
+        self.assertEqual(f.body, body)
+        
+        
+        # The same with masking given out gradually
+        mask = os.urandom(4)
+        bytes = Frame(opcode=OPCODE_TEXT, body=body, fin=1, masking_key=mask).build()
+        f = Frame()
+        f.parser.send(bytes[:10])
+        self.assertTrue(f.masking_key is None)
+        self.assertEqual(f.payload_length, 65536)
+        
+        # parse the mask gradually
+        f.parser.send(bytes[10:12])
+        f.parser.send(bytes[12:])
+        self.assertEqual(f.unmask(f.body), body)
+        
+    def test_frame_sized_126(self):
+        body = b'*'*256
+        bytes = Frame(opcode=OPCODE_TEXT, body=body, fin=1).build()
+
+        f = Frame()
+        # determine how the size is stored
+        f.parser.send(bytes[:3])
+        self.assertTrue(f.masking_key is None)
+        # that's a large frame indeed
+        self.assertEqual(f.payload_length, 126)
+
+        # this will compute the actual application data size
+        # it will also read the first byte of data
+        # indeed the length is found from byte 3 to 10
+        f.parser.send(bytes[3:11])
+        self.assertEqual(f.payload_length, 256)
+        
+        # parse the rest of our data
+        f.parser.send(bytes[11:])
+        self.assertEqual(f.body, body)
+        
+        
+        # The same but this time we provide enough
+        # bytes so that the application's data length
+        # can be computed from the first generator's send call
+        f = Frame()
+        f.parser.send(bytes[:10])
+        self.assertTrue(f.masking_key is None)
+        self.assertEqual(f.payload_length, 256)
+        
+        # parse the rest of our data
+        f.parser.send(bytes[10:])
+        self.assertEqual(f.body, body)
+        
 if __name__ == '__main__':
     suite = unittest.TestSuite()
     loader = unittest.TestLoader()
diff --git a/test/test_websocket.py b/test/test_websocket.py
@@ -4,7 +4,7 @@
 import socket
 import struct
 
-from mock import MagicMock, call
+from mock import MagicMock, call, patch
 
 from ws4py.framing import Frame, \
      OPCODE_CONTINUATION, OPCODE_TEXT, \
@@ -36,6 +36,58 @@ def test_get_underlying_connection(self):
         ws = WebSocket(sock=m)
         self.assertEqual(ws.connection, m)
 
+    def test_close_connection(self):
+        m = MagicMock()
+        ws = WebSocket(sock=m)
+        ws.close_connection()
+        m.shutdown.assert_called_once_with(socket.SHUT_RDWR)
+        m.close.assert_called_once_with()
+        self.assertIsNone(ws.connection)
+
+        m = MagicMock()
+        m.close = MagicMock(side_effect=RuntimeError)
+        ws = WebSocket(sock=m)
+        ws.close_connection()
+        self.assertIsNone(ws.connection)
+
+    def test_terminate_with_closing(self):
+        m = MagicMock()
+        s = MagicMock()
+        c = MagicMock()
+        cc = MagicMock()
+        
+        ws = WebSocket(sock=m)
+        with patch.multiple(ws, closed=c, close_connection=cc):
+            ws.stream = s
+            ws.stream.closing = CloseControlMessage(code=1000, reason='test closing')
+            ws.terminate()
+            self.assertTrue(ws.client_terminated)
+            self.assertTrue(ws.server_terminated)
+            self.assertTrue(ws.terminated)
+            c.assert_called_once_with(1000, 'test closing')
+            cc.assert_called_once_with()
+            self.assertIsNone(ws.stream)
+            self.assertIsNone(ws.environ)
+        
+    def test_terminate_without_closing(self):
+        m = MagicMock()
+        s = MagicMock()
+        c = MagicMock()
+        cc = MagicMock()
+        
+        ws = WebSocket(sock=m)
+        with patch.multiple(ws, closed=c, close_connection=cc):
+            ws.stream = s
+            ws.stream.closing = None
+            ws.terminate()
+            self.assertTrue(ws.client_terminated)
+            self.assertTrue(ws.server_terminated)
+            self.assertTrue(ws.terminated)
+            c.assert_called_once_with(1006, "Going away")
+            cc.assert_called_once_with()
+            self.assertIsNone(ws.stream)
+            self.assertIsNone(ws.environ)
+        
     def test_cannot_process_more_data_when_stream_is_terminated(self):
         m = MagicMock()
         ws = WebSocket(sock=m)
@@ -56,6 +108,67 @@ def test_no_bytes_were_read(self):
         ws = WebSocket(sock=m)
         self.assertFalse(ws.once())
 
+    def test_send_bytes_without_masking(self):
+        tm = TextMessage(b'hello world').single()
+        
+        m = MagicMock()
+        ws = WebSocket(sock=m)
+        ws.send(b'hello world')
+        m.sendall.assert_called_once_with(tm)
+        
+    def test_send_bytes_with_masking(self):
+        tm = TextMessage(b'hello world').single(mask=True)
+        
+        m = MagicMock()
+        ws = WebSocket(sock=m)
+        ws.stream = MagicMock()
+        ws.stream.always_mask = True
+        ws.stream.text_message.return_value.single.return_value = tm
+        
+        ws.send(b'hello world')
+        m.sendall.assert_called_once_with(tm)
+        
+    def test_send_message_without_masking(self):
+        tm = TextMessage(b'hello world')
+        
+        m = MagicMock()
+        ws = WebSocket(sock=m)
+        ws.send(tm)
+        m.sendall.assert_called_once_with(tm.single())
+        
+    def test_send_generator_without_masking(self):
+        tm0 = b'hello'
+        tm1 = b'world'
+        
+        def datasource():
+            yield tm0
+            yield tm1
+
+        gen = datasource()
+        
+        m = MagicMock()
+        ws = WebSocket(sock=m)
+        ws.send(gen)
+        self.assertEqual(m.sendall.call_count, 2)
+        self.assertRaises(StopIteration, next, gen)
+        
+    def test_sending_unknown_datetype(self):
+        m = MagicMock()
+        ws = WebSocket(sock=m)
+        self.assertRaises(ValueError, ws.send, 123)
+
+    def test_closing_message_received(self):
+        s = MagicMock()
+        m = MagicMock()
+        c = MagicMock()
+        
+        ws = WebSocket(sock=m)
+        with patch.multiple(ws, close=c):
+            ws.stream = s
+            ws.stream.closing = CloseControlMessage(code=1000, reason='test closing')
+            ws.process(b'unused for this test')
+            c.assert_called_once_with(1000, 'test closing')
+            
 
 if __name__ == '__main__':
     suite = unittest.TestSuite()
diff --git a/ws4py/framing.py b/ws4py/framing.py
@@ -181,6 +181,7 @@ def _parsing(self):
             some_bytes = b''
 
         if self.payload_length == 127:
+            # This will compute the actual application data size
             if len(buf) < 8:
                 nxt_buf_size = 8 - len(buf)
                 some_bytes = (yield nxt_buf_size)
@@ -191,6 +192,7 @@ def _parsing(self):
                         some_bytes = some_bytes + b
                 if len(some_bytes) > 8:
                     buf = some_bytes[8:]
+                    some_bytes = some_bytes[:8]
             else:
                 some_bytes = buf[:8]
                 buf = buf[8:]
@@ -210,6 +212,7 @@ def _parsing(self):
                         some_bytes = some_bytes + b
                 if len(some_bytes) > 2:
                     buf = some_bytes[2:]
+                    some_bytes = some_bytes[:2]
             else:
                 some_bytes = buf[:2]
                 buf = buf[2:]
