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()