def test_execute_lambda(self, mock_start_transaction, mock_session,
mock_transaction, mock_stream_cursor,
mock_buffered_cursor, mock_is_instance,
mock_executor, mock_throw_if_closed):
mock_start_transaction.return_value = mock_transaction
mock_transaction.execute_lambda.return_value = MOCK_RESULT
mock_transaction.commit.return_value = None
mock_is_instance.return_value = True
mock_stream_cursor.return_value = mock_stream_cursor
mock_buffered_cursor.return_value = MOCK_RESULT
mock_lambda = Mock()
mock_lambda.return_value = MOCK_RESULT
qldb_session = QldbSession(mock_session, MOCK_READ_AHEAD,
MOCK_RETRY_LIMIT, mock_executor)
result = qldb_session.execute_lambda(mock_lambda, None)
mock_throw_if_closed.assert_called_once_with()
mock_start_transaction.assert_called_once_with()
mock_lambda.assert_called_once()
mock_transaction.commit.assert_called_once_with()
mock_is_instance.assert_called_with(MOCK_RESULT, mock_stream_cursor)
mock_buffered_cursor.assert_called_once_with(MOCK_RESULT)
self.assertEqual(result, MOCK_RESULT)
def test_execute_lambda(self, mock_start_transaction, mock_session,
mock_transaction, mock_stream_cursor,
mock_buffered_cursor, mock_is_instance,
mock_executor):
mock_start_transaction.return_value = mock_transaction
mock_transaction.execute_lambda.return_value = MOCK_RESULT
mock_transaction._commit.return_value = None
mock_is_instance.return_value = True
mock_stream_cursor.return_value = mock_stream_cursor
mock_buffered_cursor.return_value = MOCK_RESULT
lambda_execution_context = _LambdaExecutionContext()
mock_lambda = Mock()
mock_lambda.return_value = MOCK_RESULT
qldb_session = QldbSession(mock_session, MOCK_READ_AHEAD,
mock_executor, MOCK_DRIVER_RELEASE)
result = qldb_session._execute_lambda(mock_lambda,
MOCK_DEFAULT_RETRY_CONFIG,
lambda_execution_context)
mock_start_transaction.assert_called_once_with()
mock_lambda.assert_called_once()
mock_transaction._commit.assert_called_once_with()
mock_is_instance.assert_called_with(MOCK_RESULT, mock_stream_cursor)
mock_buffered_cursor.assert_called_once_with(MOCK_RESULT)
self.assertEqual(result, MOCK_RESULT)
def test_close(self, mock_session, mock_executor):
qldb_session = QldbSession(mock_session, MOCK_READ_AHEAD,
MOCK_RETRY_LIMIT, mock_executor)
qldb_session._is_closed = False
qldb_session.close()
self.assertTrue(qldb_session._is_closed)
mock_session.close.assert_called_once_with()
def test_no_throw_abort_transaction_is_none(self, mock_session,
mock_executor):
qldb_session = QldbSession(mock_session, MOCK_READ_AHEAD,
mock_executor)
qldb_session._no_throw_abort()
mock_session._abort_transaction.assert_called_once_with()
def test_execute_lambda_unknown_exception(
self, mock_start_transaction, mock_session, mock_transaction,
mock_is_occ_conflict_exception, mock_is_retryable_exception,
mock_is_invalid_session_exception, mock_executor,
mock_no_throw_abort):
error = KeyError()
mock_start_transaction.return_value = mock_transaction
mock_transaction.transaction_id = 'transaction_id'
mock_is_occ_conflict_exception.return_value = False
mock_is_retryable_exception.return_value = False
mock_is_invalid_session_exception.return_value = False
qldb_session = QldbSession(mock_session, MOCK_READ_AHEAD,
mock_executor)
mock_lambda = Mock()
mock_lambda.side_effect = error
with self.assertRaises(ExecuteError) as cm:
qldb_session._execute_lambda(mock_lambda)
assert_execute_error(self, cm.exception, error, False, False,
mock_transaction.transaction_id)
self.assertTrue(qldb_session._is_alive)
mock_no_throw_abort.assert_called_once_with()
mock_transaction._commit.assert_not_called()
mock_transaction._close_child_cursors.assert_called_once_with()
def test_abort_or_close_when_closed(self, mock_session, mock_executor):
qldb_session = QldbSession(mock_session, MOCK_READ_AHEAD,
MOCK_RETRY_LIMIT, mock_executor)
qldb_session._is_closed = True
self.assertFalse(qldb_session._abort_or_close())
mock_session.abort_transaction.assert_not_called()
def test_execute_lambda_occ_conflict(self, mock_start_transaction,
mock_session, mock_transaction,
mock_is_occ_conflict_exception,
mock_is_retryable_exception,
mock_is_invalid_session_exception,
mock_executor, mock_no_throw_abort):
ce = ClientError(MOCK_CLIENT_ERROR_MESSAGE, MOCK_MESSAGE)
mock_start_transaction.return_value = mock_transaction
mock_transaction.transaction_id = 'transaction_id'
mock_is_retryable_exception.return_value = False
mock_is_invalid_session_exception.return_value = False
mock_is_occ_conflict_exception.return_value = True
qldb_session = QldbSession(mock_session, MOCK_READ_AHEAD,
mock_executor)
mock_lambda = Mock()
mock_lambda.side_effect = ce
with self.assertRaises(ExecuteError) as cm:
qldb_session._execute_lambda(mock_lambda)
assert_execute_error(self, cm.exception, ce, False, False,
mock_transaction.transaction_id)
mock_no_throw_abort.assert_not_called()
mock_transaction._commit.assert_not_called()
mock_transaction._close_child_cursors.assert_called_once_with()
def test_retry_sleep(self, mock_sleep, mock_random):
mock_random.return_value = 1
attempt_number = 1
exponential_back_off = min(MOCK_SLEEP_CAP_MS,
pow(MOCK_SLEEP_BASE_MS, attempt_number))
sleep_value = 1 * (exponential_back_off + 1) / 1000
QldbSession._retry_sleep(attempt_number)
mock_sleep.assert_called_once_with(sleep_value)
def test_end_session(self, mock_session, mock_executor):
qldb_session = QldbSession(mock_session, MOCK_READ_AHEAD,
mock_executor)
qldb_session._is_alive = True
qldb_session._end_session()
mock_session._close.assert_called_once_with()
self.assertFalse(qldb_session._is_alive)
def test_close_twice(self, mock_session, mock_executor):
qldb_session = QldbSession(mock_session, MOCK_READ_AHEAD,
mock_executor, MOCK_DRIVER_RELEASE)
qldb_session._is_closed = False
qldb_session._close()
qldb_session._close()
self.assertTrue(qldb_session._is_closed)
def test_throw_if_closed_when_closed(self, mock_executor, mock_session,
mock_release_session,
mock_logger_error):
qldb_session = QldbSession(mock_session, MOCK_READ_AHEAD,
mock_executor, mock_release_session)
qldb_session._is_closed = True
self.assertRaises(SessionClosedError, qldb_session._throw_if_closed)
mock_logger_error.assert_called_once()
def test_abort_or_close_client_error(self, mock_session, mock_executor):
ce = ClientError(MOCK_CLIENT_ERROR_MESSAGE, MOCK_MESSAGE)
mock_session.abort_transaction.side_effect = ce
qldb_session = QldbSession(mock_session, MOCK_READ_AHEAD,
MOCK_RETRY_LIMIT, mock_executor)
self.assertFalse(qldb_session._abort_or_close())
mock_session.abort_transaction.assert_called_once()
self.assertTrue(qldb_session._is_closed)
def test_throw_if_closed_session_closed_error(self, mock_session,
mock_logger_error,
mock_executor):
qldb_session = QldbSession(mock_session, MOCK_READ_AHEAD,
MOCK_RETRY_LIMIT, mock_executor)
qldb_session._is_closed = True
self.assertRaises(SessionClosedError, qldb_session.throw_if_closed)
mock_logger_error.assert_called_once()
def test_no_throw_abort_client_error(self, mock_session, mock_executor,
mock_transaction,
mock_logger_warning):
qldb_session = QldbSession(mock_session, MOCK_READ_AHEAD,
MOCK_RETRY_LIMIT, mock_executor)
ce = ClientError(MOCK_CLIENT_ERROR_MESSAGE, MOCK_MESSAGE)
mock_transaction.abort.side_effect = ce
qldb_session._no_throw_abort(mock_transaction)
mock_transaction.abort.assert_called_once_with()
mock_logger_warning.assert_called_once()
def test_throw_if_closed_not_closed(self, mock_executor, mock_session,
mock_release_session,
mock_logger_error):
qldb_session = QldbSession(mock_session, MOCK_READ_AHEAD,
mock_executor, mock_release_session)
qldb_session._is_closed = False
result = qldb_session._throw_if_closed()
self.assertEqual(result, None)
mock_logger_error.assert_not_called()
def test_no_throw_abort_client_error(self, mock_session,
mock_logger_warning, mock_executor):
mock_session._abort_transaction.side_effect = ClientError(
MOCK_CLIENT_ERROR_MESSAGE, 'message')
qldb_session = QldbSession(mock_session, MOCK_READ_AHEAD,
mock_executor)
qldb_session._no_throw_abort()
self.assertFalse(qldb_session._is_alive)
mock_logger_warning.assert_called_once()
def test_execute_statement(self, mock_execute_lambda, mock_session,
mock_executor):
retry_indicator = Mock()
mock_execute_lambda.return_value = mock_execute_lambda
qldb_session = QldbSession(mock_session, MOCK_READ_AHEAD,
MOCK_RETRY_LIMIT, mock_executor)
result = qldb_session.execute_statement(
MOCK_STATEMENT,
MOCK_PARAMETER_1,
MOCK_PARAMETER_2,
retry_indicator=retry_indicator)
mock_execute_lambda.assert_called_once()
self.assertEqual(result, mock_execute_lambda)
def test_return_session_with_invalid_session_exception(
self, execute_lambda_1, execute_lambda_2,
mock_lambda_context_increment_execution_attempt):
"""
The test asserts that a bad session is not returned to the pool.
We add two mock sessions to the pool. mock_session_1._execute_lambda returns an InvalidSessionException
and mock_session_2._execute_lambda succeeds.
After executing driver.execute_lambda we assert if the pool has just one session which should be
mock_session_2.
"""
mock_lambda = Mock()
session_1 = Mock()
session_2 = Mock()
driver = QldbDriver(MOCK_LEDGER_NAME)
session_1 = QldbSession(session_1, driver._read_ahead,
driver._executor, driver._release_session)
session_2 = QldbSession(session_2, driver._read_ahead,
driver._executor, driver._release_session)
mock_invalid_session_error_message = {
'Error': {
'Code': 'InvalidSessionException',
'Message': MOCK_MESSAGE
}
}
mock_invalid_session_error = ClientError(
mock_invalid_session_error_message, MOCK_MESSAGE)
execute_lambda_1.side_effect = mock_invalid_session_error
session_1._execute_lambda = execute_lambda_1
session_1._is_closed = True
execute_lambda_2.return_value = MOCK_MESSAGE
session_2._execute_lambda = execute_lambda_2
session_2._is_closed = False
# adding sessions to the driver pool
driver._pool.put(session_1)
driver._pool.put(session_2)
result = driver.execute_lambda(mock_lambda)
self.assertEqual(session_1._execute_lambda.call_count, 1)
self.assertEqual(session_1._is_closed, True)
self.assertEqual(session_2._execute_lambda.call_count, 1)
self.assertEqual(session_2._is_closed, False)
self.assertEqual(driver._pool_permits._value, 10)
self.assertEqual(driver._pool.qsize(), 1)
self.assertEqual(session_2, driver._pool.get_nowait())
self.assertEqual(
mock_lambda_context_increment_execution_attempt.call_count, 0)
self.assertEqual(result, MOCK_MESSAGE)
def test_start_transaction(self, mock_session, mock_transaction,
mock_executor):
qldb_session = QldbSession(mock_session, MOCK_READ_AHEAD,
mock_executor)
mock_transaction.return_value = mock_transaction
mock_session._start_transaction.return_value = MOCK_TRANSACTION_RESULT
transaction = qldb_session._start_transaction()
mock_session._start_transaction.assert_called_once_with()
mock_transaction.assert_called_once_with(qldb_session._session,
qldb_session._read_ahead,
MOCK_TRANSACTION_ID,
mock_executor)
self.assertEqual(transaction, mock_transaction)
def test_list_tables(self, mock_execute_statement, mock_session,
mock_executor):
mock_execute_statement.return_value = MOCK_LIST_TABLES_RESULT
qldb_session = QldbSession(mock_session, MOCK_READ_AHEAD,
MOCK_RETRY_LIMIT, mock_executor)
table_names = qldb_session.list_tables()
count = 0
for result in table_names:
self.assertEqual(result, MOCK_LIST_TABLES_RESULT_STRING[count])
count += 1
mock_execute_statement.assert_called_once_with(
"SELECT VALUE name FROM information_schema.user_tables WHERE status "
"= 'ACTIVE'")
def test_execute_lambda_invalid_session_exception(
self, mock_start_transaction, mock_session, mock_transaction,
mock_is_invalid_session, mock_executor, mock_no_throw_abort):
mock_client = Mock()
ce = ClientError(MOCK_CLIENT_ERROR_MESSAGE, 'message')
mock_start_transaction.return_value = mock_transaction
type(mock_session).client = PropertyMock(return_value=mock_client)
type(mock_session).ledger_name = PropertyMock(
return_value=MOCK_LEDGER_NAME)
mock_session._start_session = mock_session
mock_is_invalid_session.return_value = True
qldb_session = QldbSession(mock_session, MOCK_READ_AHEAD,
mock_executor, MOCK_DRIVER_RELEASE)
lambda_execution_context = _LambdaExecutionContext()
mock_lambda = Mock()
mock_lambda.side_effect = ce
execution_attempt_before = lambda_execution_context.execution_attempt
self.assertRaises(ClientError, qldb_session._execute_lambda,
mock_lambda, MOCK_RETRY_CONFIG_WITH_1_RETRY,
lambda_execution_context)
mock_start_transaction.assert_called_with()
mock_no_throw_abort.assert_not_called()
mock_transaction._commit.assert_not_called()
self.assertEqual(mock_session._start_session.call_count, 0)
self.assertEqual(mock_lambda.call_count, 1)
self.assertEqual(execution_attempt_before,
lambda_execution_context.execution_attempt)
self.assertEqual(qldb_session._is_closed, True)
def test_execute_lambda_retriable_http_exception(
self, mock_start_transaction, mock_session, mock_transaction,
mock_executor, mock_no_throw_abort):
for error_class in RETRYABLE_HTTP_ERRORS:
ce = error_class(endpoint_url="http://a-url")
mock_start_transaction.return_value = mock_transaction
qldb_session = QldbSession(mock_session, MOCK_READ_AHEAD,
mock_executor, MOCK_DRIVER_RELEASE)
lambda_execution_context = _LambdaExecutionContext()
mock_lambda = Mock()
mock_lambda.side_effect = ce
self.assertRaises(error_class, qldb_session._execute_lambda,
mock_lambda, MOCK_DEFAULT_RETRY_CONFIG,
lambda_execution_context)
mock_start_transaction.assert_has_calls(
[call(), call(), call(),
call(), call()])
mock_no_throw_abort.assert_has_calls([
call(mock_transaction),
call(mock_transaction),
call(mock_transaction),
call(mock_transaction),
call(mock_transaction)
])
self.assertEqual(mock_lambda.call_count,
MOCK_DEFAULT_RETRY_CONFIG.retry_limit + 1)
mock_transaction._commit.assert_not_called()
def test_execute_lambda_occ_conflict(self, mock_start_transaction,
mock_session, mock_transaction,
mock_is_occ_conflict_exception,
mock_is_retriable_exception,
mock_logger_warning, mock_executor,
mock_no_throw_abort):
ce = ClientError(MOCK_CLIENT_ERROR_MESSAGE, MOCK_MESSAGE)
mock_start_transaction.return_value = mock_transaction
mock_is_occ_conflict_exception.return_value = True
mock_is_retriable_exception.return_value = False
qldb_session = QldbSession(mock_session, MOCK_READ_AHEAD,
mock_executor, MOCK_DRIVER_RELEASE)
lambda_execution_context = _LambdaExecutionContext()
mock_lambda = Mock()
mock_lambda.side_effect = ce
self.assertRaises(ClientError, qldb_session._execute_lambda,
mock_lambda, MOCK_DEFAULT_RETRY_CONFIG,
lambda_execution_context)
mock_start_transaction.assert_has_calls(
[call(), call(), call(), call(),
call()])
self.assertEqual(mock_no_throw_abort.call_count, 0)
mock_is_occ_conflict_exception.assert_has_calls(
[call(ce), call(ce),
call(ce), call(ce),
call(ce)])
self.assertEqual(mock_lambda.call_count,
MOCK_DEFAULT_RETRY_CONFIG.retry_limit + 1)
self.assertEqual(mock_logger_warning.call_count,
MOCK_DEFAULT_RETRY_CONFIG.retry_limit)
mock_transaction._commit.assert_not_called()
def test_execute_lambda_start_transaction_error(self,
mock_start_transaction,
mock_session,
mock_executor,
mock_no_throw_abort):
mock_bad_request_exception_message = {
'Error': {
'Code': 'BadRequestException',
'Message': MOCK_MESSAGE
}
}
mock_bad_request_exception = ClientError(
mock_bad_request_exception_message, MOCK_MESSAGE)
ste = StartTransactionError(mock_bad_request_exception)
mock_start_transaction.side_effect = ste
lambda_execution_context = _LambdaExecutionContext()
mock_lambda = Mock()
qldb_session = QldbSession(mock_session, MOCK_READ_AHEAD,
mock_executor, MOCK_DRIVER_RELEASE)
self.assertRaises(ClientError, qldb_session._execute_lambda,
mock_lambda, MOCK_DEFAULT_RETRY_CONFIG,
lambda_execution_context)
self.assertEqual(mock_no_throw_abort.call_count,
MOCK_DEFAULT_RETRY_CONFIG.retry_limit + 1)
def test_execute_lambda_invalid_session_exception(
self, mock_start_transaction, mock_session, mock_transaction,
mock_is_invalid_session, mock_logger, mock_executor,
mock_throw_if_closed, mock_no_throw_abort):
mock_client = Mock()
ce = ClientError(MOCK_CLIENT_ERROR_MESSAGE, 'message')
mock_start_transaction.return_value = mock_transaction
type(mock_session).client = PropertyMock(return_value=mock_client)
type(mock_session).ledger_name = PropertyMock(
return_value=MOCK_LEDGER_NAME)
mock_session.start_session = mock_session
mock_is_invalid_session.return_value = True
qldb_session = QldbSession(mock_session, MOCK_READ_AHEAD, 1,
mock_executor)
mock_lambda = Mock()
mock_lambda.side_effect = ce
self.assertRaises(ClientError, qldb_session.execute_lambda,
mock_lambda)
mock_throw_if_closed.assert_called_once_with()
mock_start_transaction.assert_called_with()
mock_session.start_session.assert_called_with(MOCK_LEDGER_NAME,
mock_client)
mock_no_throw_abort.assert_called_with(mock_transaction)
mock_transaction.commit.assert_not_called()
self.assertEqual(mock_session.start_session.call_count,
qldb_session._retry_limit)
self.assertEqual(mock_lambda.call_count, qldb_session._retry_limit + 1)
self.assertEqual(mock_logger.warning.call_count,
qldb_session._retry_limit + 1)
self.assertEqual(mock_logger.info.call_count,
qldb_session._retry_limit)
def test_constructor(self, mock_session, mock_executor):
qldb_session = QldbSession(mock_session, MOCK_READ_AHEAD,
mock_executor)
self.assertEqual(qldb_session._is_alive, True)
self.assertEqual(qldb_session._read_ahead, MOCK_READ_AHEAD)
self.assertEqual(qldb_session._session, mock_session)
self.assertEqual(qldb_session._executor, mock_executor)
def test_constructor(self, mock_session, mock_executor):
qldb_session = QldbSession(mock_session, MOCK_READ_AHEAD,
MOCK_RETRY_LIMIT, mock_executor)
self.assertEqual(qldb_session._is_closed, False)
self.assertEqual(qldb_session._read_ahead, MOCK_READ_AHEAD)
self.assertEqual(qldb_session._retry_limit, MOCK_RETRY_LIMIT)
self.assertEqual(qldb_session._session, mock_session)
self.assertEqual(qldb_session._executor, mock_executor)
def test_constructor(self, mock_session, mock_executor,
mock_release_session):
qldb_session = QldbSession(mock_session, MOCK_READ_AHEAD,
mock_executor, mock_release_session)
self.assertEqual(qldb_session._is_closed, False)
self.assertEqual(qldb_session._read_ahead, MOCK_READ_AHEAD)
self.assertEqual(qldb_session._session, mock_session)
self.assertEqual(qldb_session._executor, mock_executor)
self.assertEqual(qldb_session._return_session_to_pool,
mock_release_session)
def test_execute_lambda_lambda_aborted_error(
self, mock_start_transaction, mock_session, mock_transaction,
mock_is_occ_conflict, mock_is_retriable_exception, mock_executor,
mock_no_throw_abort):
mock_start_transaction.return_value = mock_transaction
mock_is_occ_conflict.return_value = True
mock_is_retriable_exception.return_value = False
retry_indicator = Mock()
qldb_session = QldbSession(mock_session, MOCK_READ_AHEAD,
MOCK_RETRY_LIMIT, mock_executor)
mock_lambda = Mock()
mock_lambda.side_effect = LambdaAbortedError
self.assertRaises(LambdaAbortedError, qldb_session.execute_lambda,
mock_lambda, retry_indicator)
mock_no_throw_abort.assert_called_once_with(mock_transaction)
def test_execute_lambda_lambda_aborted_error(
self, mock_start_transaction, mock_session, mock_transaction,
mock_is_occ_conflict, mock_is_retriable_exception, mock_executor,
mock_no_throw_abort):
mock_start_transaction.return_value = mock_transaction
mock_is_occ_conflict.return_value = True
mock_is_retriable_exception.return_value = False
retry_indicator = Mock()
qldb_session = QldbSession(mock_session, MOCK_READ_AHEAD,
mock_executor, MOCK_DRIVER_RELEASE)
lambda_execution_context = _LambdaExecutionContext()
mock_lambda = Mock()
mock_lambda.side_effect = LambdaAbortedError
self.assertRaises(LambdaAbortedError, qldb_session._execute_lambda,
mock_lambda, MOCK_DEFAULT_RETRY_CONFIG,
lambda_execution_context)
mock_no_throw_abort.assert_called_once_with(mock_transaction)
