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