def test_run_with_failure(self): """ Test that the correct logging and error handling occurs when the thread encounters an error with the message handling loop. """ kmip_session = session.KmipSession(None, None, 'name') kmip_session._logger = mock.MagicMock() kmip_session._connection = mock.MagicMock() test_exception = Exception("test") kmip_session._handle_message_loop = mock.MagicMock( side_effect=[test_exception, exceptions.ConnectionClosed()]) kmip_session.run() kmip_session._logger.info.assert_any_call("Starting session: name") self.assertTrue(kmip_session._handle_message_loop.called) kmip_session._logger.info.assert_any_call( "Failure handling message loop") kmip_session._logger.exception.assert_called_once_with(test_exception) kmip_session._connection.shutdown.assert_called_once_with( socket.SHUT_RDWR) kmip_session._connection.close.assert_called_once_with() kmip_session._logger.info.assert_called_with("Stopping session: name")
def test_run(self): """ Test that the message handling loop is handled properly on normal execution. """ kmip_session = session.KmipSession(None, None, 'name') kmip_session._logger = mock.MagicMock() kmip_session._handle_message_loop = mock.MagicMock( side_effect=[None, exceptions.ConnectionClosed()]) kmip_session._connection = mock.MagicMock() kmip_session.run() kmip_session._logger.info.assert_any_call("Starting session: name") self.assertTrue(kmip_session._handle_message_loop.called) kmip_session._connection.shutdown.assert_called_once_with( socket.SHUT_RDWR) kmip_session._connection.close.assert_called_once_with() kmip_session._logger.info.assert_called_with("Stopping session: name")
def _receive_bytes(self, message_size): bytes_received = 0 message = b'' while bytes_received < message_size: partial_message = self._connection.recv( min(message_size - bytes_received, self._max_buffer_size)) if partial_message is None: break elif len(partial_message) == 0: raise exceptions.ConnectionClosed() else: bytes_received += len(partial_message) message += partial_message if bytes_received != message_size: raise ValueError( "Invalid KMIP message received. Actual message length " "does not match the advertised header length.") else: return message