On a closed connection, return buffered input before raising an excep…

…tion (Fixes #30)

src/simple_websocket/ws.py

@@ -110,9 +110,12 @@ def receive(self, timeout=None):
             if not self.event.wait(timeout=timeout):
                 return None
             self.event.clear()
+        try:
+            return self.input_buffer.pop(0)
+        except IndexError:
+            pass
         if not self.connected:  # pragma: no cover
             raise ConnectionClosed(self.close_reason, self.close_message)
-        return self.input_buffer.pop(0)
 
     def close(self, reason=None, message=None):
         """Close the WebSocket connection.

tests/test_simple_websocket_client.py

@@ -113,6 +113,19 @@ def test_receive(self, mock_wsconn, mock_socket):
         assert client.receive() == b'hello'
         assert client.receive(timeout=0) is None
 
+    @mock.patch('simple_websocket.ws.socket.socket')
+    @mock.patch('simple_websocket.ws.WSConnection')
+    def test_receive_after_close(self, mock_wsconn, mock_socket):
+        mock_socket.return_value.recv.return_value = b'x'
+        client = self.get_client(mock_wsconn, 'ws://example.com/ws', events=[
+            [TextMessage('hello')],
+        ])
+        while client.connected:
+            time.sleep(0.01)
+        assert client.receive() == 'hello'
+        with pytest.raises(simple_websocket.ConnectionClosed):
+            client.receive()
+
     @mock.patch('simple_websocket.ws.socket.socket')
     @mock.patch('simple_websocket.ws.WSConnection')
     def test_receive_ping(self, mock_wsconn, mock_socket):

tests/test_simple_websocket_server.py

@@ -113,6 +113,21 @@ def test_receive(self, mock_wsconn):
         assert server.receive() == b'hello'
         assert server.receive(timeout=0) is None
 
+    @mock.patch('simple_websocket.ws.WSConnection')
+    def test_receive_after_close(self, mock_wsconn):
+        mock_socket = mock.MagicMock()
+        mock_socket.recv.return_value = b'x'
+        server = self.get_server(mock_wsconn, {
+            'werkzeug.socket': mock_socket,
+        }, events=[
+            [TextMessage('hello')],
+        ])
+        while server.connected:
+            time.sleep(0.01)
+        assert server.receive() == 'hello'
+        with pytest.raises(simple_websocket.ConnectionClosed):
+            server.receive()
+
     @mock.patch('simple_websocket.ws.WSConnection')
     def test_receive_split_messages(self, mock_wsconn):
         mock_socket = mock.MagicMock()