def test_execute_lambda_retryable_error_and_under_retry_limit(
self, mock_get_session, mock_client, mock_retry_sleep,
mock_session, mock_release_session):
mock_client.return_value = mock_client
mock_lambda = Mock()
mock_get_session.return_value = mock_session
inner_error = Exception()
retryable_execute_error = ExecuteError(inner_error, True, False,
DEFAULT_TRANSACTION_ID)
mock_session._execute_lambda.side_effect = [
retryable_execute_error, MOCK_MESSAGE
]
mock_release_session.return_value = True
driver = QldbDriver(MOCK_LEDGER_NAME)
result = driver.execute_lambda(mock_lambda)
self.assertEqual(result, MOCK_MESSAGE)
mock_get_session.assert_has_calls([call(False), call(False)])
mock_release_session.assert_has_calls(
[call(mock_session), call(mock_session)])
mock_retry_sleep.assert_called_once_with(driver._retry_config, 1,
inner_error,
DEFAULT_TRANSACTION_ID)
self.assertEqual(mock_session._execute_lambda.call_count, 2)
def test_get_session_when_closed(self, mock_client):
mock_client.return_value = mock_client
mock_client.max_pool_connections = DEFAULT_MAX_CONCURRENT_TRANSACTIONS
qldb_driver = QldbDriver(MOCK_LEDGER_NAME)
qldb_driver._is_closed = True
self.assertRaises(DriverClosedError, qldb_driver._get_session)
def test_execute_lambda_with_InvalidSessionException(
self, mock_get_session, mock_client):
"""
The test asserts that if an InvalidSessionException is thrown by session.execute_lambda, the
driver calls _get_session
"""
mock_client.return_value = mock_client
mock_lambda = Mock()
mock_session = mock_get_session.return_value.__enter__.return_value
mock_invalid_session_error_message = {
'Error': {
'Code': 'InvalidSessionException',
'Message': MOCK_MESSAGE
}
}
mock_invalid_session_error = ClientError(
mock_invalid_session_error_message, MOCK_MESSAGE)
mock_session._execute_lambda.side_effect = [
mock_invalid_session_error, mock_invalid_session_error,
MOCK_MESSAGE
]
driver = QldbDriver(MOCK_LEDGER_NAME)
result = driver.execute_lambda(mock_lambda, mock_lambda)
self.assertEqual(mock_get_session.call_count, 3)
self.assertEqual(result, MOCK_MESSAGE)
def test_execute_lambda_when_driver_is_closed(self, mock_client):
mock_client.return_value = mock_client
mock_lambda = Mock()
driver = QldbDriver(MOCK_LEDGER_NAME)
driver._is_closed = True
self.assertRaises(DriverClosedError, driver.execute_lambda,
mock_lambda)
def test_release_session_for_none_session(self, mock_client, mock_queue,
mock_bounded_semaphore,
mock_atomic_integer):
mock_client.return_value = mock_client
qldb_driver = QldbDriver(MOCK_LEDGER_NAME)
self.assertFalse(qldb_driver._release_session(None))
mock_queue.put.assert_not_called()
mock_bounded_semaphore().release.assert_not_called()
mock_atomic_integer().increment.assert_not_called()
def poll_for_table_creation(ledger_name):
driver = QldbDriver(ledger_name)
max_poll_time = time() + 15
while True:
tables = driver.list_tables()
count = len(list(tables))
if count == 4 or time() > max_poll_time:
break
sleep(3)
def update_patient_encounter(patient_encounter: dict):
def insert_documents(transaction_executor, payload: dict):
transaction_executor.execute_statement(
"INSERT INTO PatientEncounter ?", payload)
patient_encounter["action"] = "UPDATE"
qldb_driver = QldbDriver(ledger_name=LEDGER_NAME, region_name="us-west-2")
# pylint: disable=W0108
qldb_driver.execute_lambda(
lambda x: insert_documents(x, patient_encounter))
def test_list_tables(self, mock_execute_lambda, mock_client):
mock_client.return_value = mock_client
mock_execute_lambda.return_value = MOCK_LIST_TABLES_RESULT
driver = QldbDriver(MOCK_LEDGER_NAME)
table_names = driver.list_tables()
count = 0
for result in table_names:
self.assertEqual(result, MOCK_LIST_TABLES_RESULT[count])
count += 1
def test_close(self, mock_client, mock_qldb_session1, mock_qldb_session2):
mock_client.return_value = mock_client
mock_client.max_pool_connections = DEFAULT_MAX_CONCURRENT_TRANSACTIONS
qldb_driver = QldbDriver(MOCK_LEDGER_NAME)
qldb_driver._pool = Queue()
qldb_driver._pool.put(mock_qldb_session1)
qldb_driver._pool.put(mock_qldb_session2)
qldb_driver.close()
mock_qldb_session1._end_session.assert_called_once_with()
mock_qldb_session2._end_session.assert_called_once_with()
def test_execute_lambda(self, mock_get_session, mock_client, mock_session,
mock_release_session):
mock_client.return_value = mock_client
mock_lambda = Mock()
mock_get_session.return_value = mock_session
mock_session._execute_lambda.return_value = MOCK_MESSAGE
driver = QldbDriver(MOCK_LEDGER_NAME)
result = driver.execute_lambda(mock_lambda)
mock_release_session.assert_called_once_with(mock_session)
mock_get_session.assert_called_once_with(False)
mock_session._execute_lambda.assert_called_once_with(mock_lambda)
self.assertEqual(result, MOCK_MESSAGE)
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_create_new_session(self, mock_client, mock_session_start_session,
mock_qldb_session):
mock_session_start_session.return_value = mock_session_start_session
mock_qldb_session.return_value = mock_qldb_session
mock_client.return_value = mock_client
mock_client.max_pool_connections = DEFAULT_MAX_CONCURRENT_TRANSACTIONS
qldb_driver = QldbDriver(MOCK_LEDGER_NAME)
session = qldb_driver._create_new_session()
mock_session_start_session.assert_called_once_with(
MOCK_LEDGER_NAME, qldb_driver._client)
mock_qldb_session.assert_called_once_with(mock_session_start_session,
qldb_driver._read_ahead,
qldb_driver._executor)
self.assertEqual(session, mock_qldb_session)
def test_release_session_for_closed_session(self, mock_client,
mock_logger_debug,
mock_qldb_session, mock_queue,
mock_bounded_semaphore,
mock_atomic_integer):
mock_client.return_value = mock_client
mock_client.max_pool_connections = DEFAULT_MAX_CONCURRENT_TRANSACTIONS
qldb_driver = QldbDriver(MOCK_LEDGER_NAME)
mock_qldb_session._is_closed = True
qldb_driver._release_session(mock_qldb_session)
mock_queue().put.assert_not_called()
mock_bounded_semaphore().release.assert_called_once_with()
mock_atomic_integer().increment.assert_called_once_with()
mock_logger_debug.assert_called_once()
def test_execute_lambda(self, mock_get_session, mock_client,
mock_lambda_context):
mock_client.return_value = mock_client
mock_lambda = Mock()
mock_session = mock_get_session.return_value.__enter__.return_value
mock_session._execute_lambda.return_value = MOCK_MESSAGE
mock_lambda_context.return_value = mock_lambda_context
driver = QldbDriver(MOCK_LEDGER_NAME)
result = driver.execute_lambda(mock_lambda)
mock_get_session.assert_called_once_with()
mock_session._execute_lambda.assert_called_once_with(
mock_lambda, driver._retry_config, mock_lambda_context)
self.assertEqual(result, MOCK_MESSAGE)
def get_all_patients():
"""
Retrieve all patient data
"""
def read_documents(transaction_executor):
"""
Internal function used to retrieve all Patient data.
"""
cursor = transaction_executor.execute_statement(
"SELECT * FROM Patient")
return cursor
qldb_driver = QldbDriver(ledger_name=LEDGER_NAME, region_name="us-west-2")
# pylint: disable=W0108
qldb_driver.execute_lambda(lambda x: read_documents(x))
def test_get_session_exception(self, mock_client, mock_create_new_session,
mock_bounded_semaphore,
mock_atomic_integer):
mock_client.return_value = mock_client
mock_client.max_pool_connections = DEFAULT_MAX_CONCURRENT_TRANSACTIONS
error = KeyError()
mock_create_new_session.side_effect = error
qldb_driver = QldbDriver(MOCK_LEDGER_NAME)
with self.assertRaises(ExecuteError) as cm:
qldb_driver._get_session(False)
assert_execute_error(self, cm.exception, error, True, True, None)
mock_bounded_semaphore().release.assert_called_once_with()
mock_atomic_integer().increment.assert_called_once_with()
def create_new_patient_encounter(patient_encounter: dict):
"""
Create a new, blank patient record.
"""
def insert_documents(transaction_executor, payload: dict):
"""
Internal function handling insertion of new patients.
"""
transaction_executor.execute_statement(
"INSERT INTO PatientEncounter ?", payload)
patient_encounter["action"] = "INSERT"
qldb_driver = QldbDriver(ledger_name=LEDGER_NAME, region_name="us-west-2")
# pylint: disable=W0108
qldb_driver.execute_lambda(
lambda x: insert_documents(x, patient_encounter))
def test_get_session_new_session(self, mock_client, mock_qldb_session,
mock_bounded_semaphore,
mock_atomic_integer,
mock_create_new_session):
mock_client.return_value = mock_client
mock_client.max_pool_connections = DEFAULT_MAX_CONCURRENT_TRANSACTIONS
mock_create_new_session.return_value = mock_qldb_session
qldb_driver = QldbDriver(MOCK_LEDGER_NAME)
session = qldb_driver._get_session(True)
mock_bounded_semaphore().acquire.assert_called_once_with(
timeout=DEFAULT_TIMEOUT_SECONDS)
mock_atomic_integer().decrement.assert_called_once_with()
mock_create_new_session.assert_called_once_with()
self.assertEqual(session, mock_qldb_session)
def test_constructor_with_max_concurrent_transactions_less_than_client_max_concurrent_transactions(
self, mock_client, mock_bounded_semaphore, mock_atomic_integer,
mock_queue):
mock_queue.return_value = mock_queue
mock_atomic_integer.return_value = mock_atomic_integer
mock_bounded_semaphore.return_value = mock_bounded_semaphore
mock_client.return_value = mock_client
mock_client.max_pool_connections = DEFAULT_MAX_CONCURRENT_TRANSACTIONS
new_max_concurrent_transactions = DEFAULT_MAX_CONCURRENT_TRANSACTIONS - 1
qldb_driver = QldbDriver(
MOCK_LEDGER_NAME,
max_concurrent_transactions=new_max_concurrent_transactions)
self.assertEqual(qldb_driver._ledger_name, MOCK_LEDGER_NAME)
self.assertEqual(qldb_driver._retry_config.retry_limit,
DEFAULT_RETRY_LIMIT)
self.assertEqual(qldb_driver._retry_config.base, DEFAULT_BACKOFF_BASE)
self.assertEqual(qldb_driver._read_ahead, DEFAULT_READ_AHEAD)
self.assertEqual(qldb_driver._pool_permits, mock_bounded_semaphore)
self.assertEqual(qldb_driver._pool_permits_counter,
mock_atomic_integer)
self.assertEqual(qldb_driver._pool, mock_queue)
mock_bounded_semaphore.assert_called_once_with(
new_max_concurrent_transactions)
mock_atomic_integer.assert_called_once_with(
new_max_concurrent_transactions)
mock_queue.assert_called_once_with()
def update_patient(patient: dict):
"""
Update a patient with the provided dataset.
"""
def update_documents(transaction_executor, payload: dict):
"""
Internal function for updating a patient record.
"""
transaction_executor.execute_statement("INSERT INTO Patient ?",
payload)
patient["action"] = "UPDATE"
patient["date_of_birth"] = str(patient["date_of_birth"])
qldb_driver = QldbDriver(ledger_name=LEDGER_NAME, region_name="us-west-2")
# pylint: disable=W0108
qldb_driver.execute_lambda(lambda x: update_documents(x, patient))
def create_qldb_driver(ledger_name=Constants.LEDGER_NAME,
region_name=None,
endpoint_url=None,
boto3_session=None):
"""
Create a QLDB driver for executing transactions.
:type ledger_name: str
:param ledger_name: The QLDB ledger name.
:type region_name: str
:param region_name: See [1].
:type endpoint_url: str
:param endpoint_url: See [1].
:type boto3_session: :py:class:`boto3.session.Session`
:param boto3_session: The boto3 session to create the client with (see [1]).
:rtype: :py:class:`pyqldb.driver.qldb_driver.QldbDriver`
:return: A QLDB driver object.
[1]: `Boto3 Session.client Reference <https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html#boto3.session.Session.client>`.
"""
qldb_driver = QldbDriver(ledger_name=ledger_name,
region_name=region_name,
endpoint_url=endpoint_url,
boto3_session=boto3_session)
return qldb_driver
def test_constructor_with_valid_config(self, mock_config_merge,
mock_client, mock_bounded_semaphore,
mock_atomic_integer, mock_queue):
mock_queue.return_value = mock_queue
mock_atomic_integer.return_value = mock_atomic_integer
mock_bounded_semaphore.return_value = mock_bounded_semaphore
mock_client.return_value = mock_client
mock_config_merge.return_value = mock_config_merge
mock_config_merge.max_pool_connections = DEFAULT_MAX_CONCURRENT_TRANSACTIONS
qldb_driver = QldbDriver(MOCK_LEDGER_NAME, config=MOCK_CONFIG)
mock_config_merge.assert_called_once()
mock_client.assert_called_once_with(DEFAULT_SESSION_NAME,
aws_access_key_id=None,
aws_secret_access_key=None,
aws_session_token=None,
config=mock_config_merge,
endpoint_url=None,
region_name=None,
verify=None)
self.assertEqual(qldb_driver._ledger_name, MOCK_LEDGER_NAME)
self.assertEqual(qldb_driver._retry_config.retry_limit,
DEFAULT_RETRY_LIMIT)
self.assertEqual(qldb_driver._retry_config.base, DEFAULT_BACKOFF_BASE)
self.assertEqual(qldb_driver._read_ahead, DEFAULT_READ_AHEAD)
self.assertEqual(qldb_driver._pool_permits, mock_bounded_semaphore)
self.assertEqual(qldb_driver._pool_permits_counter,
mock_atomic_integer)
self.assertEqual(qldb_driver._pool, mock_queue)
mock_bounded_semaphore.assert_called_once_with(
DEFAULT_MAX_CONCURRENT_TRANSACTIONS)
mock_atomic_integer.assert_called_once_with(
DEFAULT_MAX_CONCURRENT_TRANSACTIONS)
mock_queue.assert_called_once_with()
def test_execute_lambda_retryable_error_and_exceed_retry_limit(
self, mock_get_session, mock_client, mock_retry_sleep,
mock_session, mock_release_session):
mock_client.return_value = mock_client
mock_lambda = Mock()
mock_get_session.return_value = mock_session
inner_error = Exception()
retryable_execute_error = ExecuteError(inner_error, True, False,
DEFAULT_TRANSACTION_ID)
mock_session._execute_lambda.side_effect = [
retryable_execute_error, retryable_execute_error,
retryable_execute_error
]
mock_release_session.return_value = True
retryConfig = RetryConfig(retry_limit=2)
driver = QldbDriver(MOCK_LEDGER_NAME, retry_config=retryConfig)
self.assertRaises(Exception, driver.execute_lambda, mock_lambda)
mock_get_session.assert_has_calls(
[call(False), call(False), call(False)])
mock_release_session.assert_has_calls(
[call(mock_session),
call(mock_session),
call(mock_session)])
mock_retry_sleep.assert_has_calls([
call(driver._retry_config, 1, inner_error, DEFAULT_TRANSACTION_ID),
call(driver._retry_config, 2, inner_error, DEFAULT_TRANSACTION_ID)
])
self.assertEqual(mock_session._execute_lambda.call_count, 3)
def test_context_manager(self, mock_client, mock_close):
mock_client.return_value = mock_client
mock_client.max_pool_connections = DEFAULT_MAX_CONCURRENT_TRANSACTIONS
with QldbDriver(MOCK_LEDGER_NAME):
pass
mock_close.assert_called_once_with()
def test_get_session_existing_session(
self, mock_client, mock_session_start_session, mock_qldb_session,
mock_release_session, mock_bounded_semaphore, mock_atomic_integer,
mock_logger_debug):
mock_session_start_session.return_value = mock_session_start_session
mock_qldb_session.return_value = mock_qldb_session
mock_client.return_value = mock_client
mock_client.max_pool_connections = DEFAULT_MAX_CONCURRENT_TRANSACTIONS
qldb_driver = QldbDriver(MOCK_LEDGER_NAME)
qldb_driver._pool = Queue()
qldb_driver._pool.put(mock_qldb_session)
session = qldb_driver._get_session()
self.assertEqual(session, mock_qldb_session)
mock_bounded_semaphore().acquire.assert_called_once_with(
timeout=DEFAULT_TIMEOUT_SECONDS)
mock_atomic_integer().decrement.assert_called_once_with()
self.assertEqual(mock_logger_debug.call_count, 2)
def test_constructor_with_boto3_session(self, mock_config_merge):
mock_session = Mock(spec=MOCK_BOTO3_SESSION)
mock_config_merge.return_value = mock_config_merge
mock_config_merge.max_pool_connections = DEFAULT_MAX_CONCURRENT_TRANSACTIONS
qldb_driver = QldbDriver(MOCK_LEDGER_NAME,
boto3_session=mock_session,
config=MOCK_CONFIG)
mock_session.client.assert_called_once_with(DEFAULT_SESSION_NAME,
config=mock_config_merge,
endpoint_url=None,
verify=None)
self.assertEqual(qldb_driver._client, mock_session.client())
def test_constructor_with_boto3_session(self):
mock_session = Mock(spec=MOCK_BOTO3_SESSION)
qldb_driver = QldbDriver(MOCK_LEDGER_NAME,
boto3_session=mock_session,
config=MOCK_CONFIG)
mock_session.client.assert_called_once_with(DEFAULT_SESSION_NAME,
config=MOCK_CONFIG,
endpoint_url=None,
verify=None)
self.assertEqual(qldb_driver._client, mock_session.client())
self.assertTrue(
SERVICE_DESCRIPTION in qldb_driver._config.user_agent_extra)
def test_execute_lambda_invalid_session_exception_and_0_retry_limit(
self, mock_get_session, mock_client, mock_retry_sleep,
mock_session, mock_release_session):
mock_client.return_value = mock_client
mock_lambda = Mock()
mock_get_session.return_value = mock_session
invalid_session_exception = ExecuteError(Exception(), True, True)
mock_session._execute_lambda.side_effect = [
invalid_session_exception, MOCK_MESSAGE
]
mock_release_session.return_value = True
retryConfig = RetryConfig(retry_limit=0)
driver = QldbDriver(MOCK_LEDGER_NAME, retry_config=retryConfig)
result = driver.execute_lambda(mock_lambda)
self.assertEqual(result, MOCK_MESSAGE)
mock_get_session.assert_has_calls([call(False), call(False)])
mock_release_session.assert_has_calls(
[call(mock_session), call(mock_session)])
mock_retry_sleep.assert_not_called()
self.assertEqual(mock_session._execute_lambda.call_count, 2)
def test_get_session_session_pool_empty_error(
self, mock_client, mock_bounded_semaphore,
mock_session_pool_empty_error, mock_logger_debug):
mock_client.return_value = mock_client
mock_client.max_pool_connections = DEFAULT_MAX_CONCURRENT_TRANSACTIONS
mock_bounded_semaphore().acquire.return_value = False
mock_session_pool_empty_error.return_value = Exception
qldb_driver = QldbDriver(MOCK_LEDGER_NAME)
self.assertRaises(Exception, qldb_driver._get_session)
mock_session_pool_empty_error.assert_called_once_with(
DEFAULT_TIMEOUT_SECONDS)
mock_logger_debug.assert_called_once()
def test_get_session_exception(self, mock_client, mock_qldb_session,
mock_bounded_semaphore,
mock_atomic_integer):
mock_client.return_value = mock_client
mock_client.max_pool_connections = DEFAULT_MAX_CONCURRENT_TRANSACTIONS
qldb_driver = QldbDriver(MOCK_LEDGER_NAME)
mock_qldb_session.side_effect = Exception
self.assertRaises(Exception, qldb_driver._get_session)
mock_bounded_semaphore().acquire.assert_called_once_with(
timeout=DEFAULT_TIMEOUT_SECONDS)
mock_atomic_integer().decrement.assert_called_once_with()
mock_bounded_semaphore().release.assert_called_once_with()
mock_atomic_integer().increment.assert_called_once_with()
