def test_next_verify_are_there_more_results_and_next_page_called(
self, mock_are_there_more_results, mock_next_page, mock_session,
mock_value_holder_to_ion_value, mock_executor):
updated_result = '1'
def next_page():
read_ahead_cursor._page = {
'NextPageToken': None,
'Values': [updated_result]
}
read_ahead_cursor._index = 0
mock_are_there_more_results.return_value = True
mock_value_holder_to_ion_value.side_effect = lambda val: val
mock_session.return_value = None
mock_next_page.return_value = None
mock_next_page.side_effect = next_page
read_ahead_cursor = ReadAheadCursor(MOCK_STATEMENT_RESULT,
mock_session, MOCK_TRANSACTION_ID,
MOCK_READ_AHEAD, mock_executor)
read_ahead_cursor._index = len(MOCK_ION_BINARY_VALUES)
self.assertEqual(next(read_ahead_cursor), updated_result)
mock_are_there_more_results.assert_called_once_with()
mock_next_page.assert_called_once_with()
mock_value_holder_to_ion_value.assert_called_once_with(updated_result)
def test_next_page(self, mock_session, mock_executor):
mock_page1 = {
'Values': [{
'IonBinary': 1
}, {
'IonBinary': 2
}],
'NextPageToken': 'token'
}
mock_page2 = {
'Values': [{
'IonBinary': 2
}, {
'IonBinary': 3
}],
'NextPageToken': None
}
mock_statement_result1 = {'Page': mock_page1}
mock_statement_result2 = {'Page': mock_page2}
mock_session.return_value = None
mock_session._fetch_page.return_value = mock_statement_result2
read_ahead_cursor = ReadAheadCursor(mock_statement_result1,
mock_session, MOCK_TRANSACTION_ID,
MOCK_READ_AHEAD, mock_executor)
read_ahead_cursor._queue = Queue()
read_ahead_cursor._queue.put(mock_statement_result2)
read_ahead_cursor._next_page()
self.assertEqual(read_ahead_cursor._page, mock_page2)
self.assertEqual(read_ahead_cursor._index, 0)
def test_close(self, mock_session, mock_executor):
mock_session.return_value = None
read_ahead_cursor = ReadAheadCursor(MOCK_STATEMENT_RESULT,
mock_session, MOCK_TRANSACTION_ID,
MOCK_READ_AHEAD, mock_executor)
read_ahead_cursor.close()
self.assertFalse(read_ahead_cursor._is_open)
def test_next_page_client_error(self, mock_session, mock_executor):
mock_session.return_value = None
mock_session._fetch_page.return_value = {'Page': MOCK_STATEMENT_RESULT}
read_ahead_cursor = ReadAheadCursor(MOCK_STATEMENT_RESULT,
mock_session, MOCK_TRANSACTION_ID,
MOCK_READ_AHEAD, mock_executor)
read_ahead_cursor._queue = Queue()
read_ahead_cursor._queue.put(
ClientError(MOCK_CLIENT_ERROR_MESSAGE, MOCK_MESSAGE))
self.assertRaises(ClientError, read_ahead_cursor._next_page)
def test_next_with_next_page_returns_empty_values_and_not_none_token(
self, mock_session, mock_executor):
read_ahead_cursor = ReadAheadCursor(MOCK_STATEMENT_RESULT,
mock_session, MOCK_TRANSACTION_ID,
MOCK_READ_AHEAD, mock_executor)
read_ahead_cursor._queue = Queue()
read_ahead_cursor._queue.put({'NextPageToken': 'token', 'Values': []})
read_ahead_cursor._queue.put({'NextPageToken': None, 'Values': []})
read_ahead_cursor._index = len(MOCK_ION_BINARY_VALUES)
self.assertRaises(StopIteration, next, read_ahead_cursor)
def test_next_page(self, mock_session, mock_executor):
mock_session.return_value = None
mock_session._fetch_page.return_value = {'Page': MOCK_STATEMENT_RESULT}
read_ahead_cursor = ReadAheadCursor(MOCK_STATEMENT_RESULT,
mock_session, MOCK_TRANSACTION_ID,
MOCK_READ_AHEAD, mock_executor)
read_ahead_cursor._queue = Queue()
read_ahead_cursor._queue.put(MOCK_STATEMENT_RESULT)
read_ahead_cursor._next_page()
self.assertEqual(read_ahead_cursor._page, MOCK_STATEMENT_RESULT)
self.assertEqual(read_ahead_cursor._index, 0)
def test_next_when_next_page_returns_empty_values_and_none_token(
self, mock_session, mock_next_page, mock_executor):
mock_session.return_value = None
def next_page():
read_ahead_cursor._page = {'NextPageToken': None, 'Values': []}
read_ahead_cursor._index = 0
read_ahead_cursor = ReadAheadCursor(MOCK_STATEMENT_RESULT,
mock_session, MOCK_TRANSACTION_ID,
MOCK_READ_AHEAD, mock_executor)
read_ahead_cursor._index = len(MOCK_ION_BINARY_VALUES)
mock_next_page.side_effect = next_page
self.assertRaises(StopIteration, next, read_ahead_cursor)
def test_populate_queue(self, mock_queue, mock_session, mock_executor):
mock_session.return_value = None
mock_queue.return_value = mock_queue
mock_page = {'NextPageToken': None, 'Values': []}
mock_session._fetch_page.return_value = {'Page': mock_page}
mock_queue.full.return_value = False
read_ahead_cursor = ReadAheadCursor(MOCK_STATEMENT_RESULT,
mock_session, MOCK_TRANSACTION_ID,
MOCK_READ_AHEAD, mock_executor)
read_ahead_cursor._queue = mock_queue
read_ahead_cursor._populate_queue()
mock_session._fetch_page.assert_called_once_with(
MOCK_TRANSACTION_ID, MOCK_STATEMENT_RESULT.get('NextPageToken'))
mock_queue.put.assert_called_once_with(mock_page, timeout=0.05)
def test_ReadAheadCursor_without_executor(self, mock_thread, mock_queue,
mock_session):
mock_session.return_value = None
mock_thread.return_value = mock_thread
mock_queue.return_value = mock_queue
read_ahead_cursor = ReadAheadCursor(MOCK_STATEMENT_RESULT,
mock_session, MOCK_TRANSACTION_ID,
MOCK_READ_AHEAD, None)
self.assertEqual(read_ahead_cursor._page, MOCK_PAGE_WITH_TOKEN)
self.assertEqual(read_ahead_cursor._session, mock_session)
self.assertEqual(read_ahead_cursor._read_ios, MOCK_READ_IOS)
self.assertEqual(read_ahead_cursor._write_ios, MOCK_WRITE_IOS)
self.assertEqual(read_ahead_cursor._processing_time_milliseconds,
MOCK_PROCESSING_TIME)
self.assertEqual(read_ahead_cursor._transaction_id,
MOCK_TRANSACTION_ID)
self.assertEqual(read_ahead_cursor._index, 0)
self.assertEqual(read_ahead_cursor._queue, mock_queue)
self.assertEqual(read_ahead_cursor._is_open, True)
mock_queue.assert_called_once_with(MOCK_READ_AHEAD - 1)
mock_thread.assert_called_once_with(
target=read_ahead_cursor._populate_queue)
mock_thread().setDaemon.assert_called_once_with(True)
mock_thread().start.assert_called_once_with()
def test_populate_queue_client_error(self, mock_session, mock_logger_debug,
mock_executor):
mock_logger_debug.return_value = None
mock_session.return_value = None
mock_session._fetch_page.side_effect = ClientError(
MOCK_CLIENT_ERROR_MESSAGE, MOCK_MESSAGE)
read_ahead_cursor = ReadAheadCursor(MOCK_STATEMENT_RESULT,
mock_session, MOCK_TRANSACTION_ID,
MOCK_READ_AHEAD, mock_executor)
read_ahead_cursor._queue = Queue(1)
read_ahead_cursor._queue.put('value to be removed')
read_ahead_cursor._populate_queue()
mock_logger_debug.assert_called_once()
self.assertIsInstance(read_ahead_cursor._queue.get(), ClientError)
self.assertEqual(read_ahead_cursor._queue.qsize(), 0)
def test_populate_queue(self, mock_queue, mock_session, mock_executor):
mock_session.return_value = None
mock_queue.return_value = mock_queue
mock_session._fetch_page.return_value = MOCK_FETCH_PAGE_RESULT_WITHOUT_TOKEN
mock_queue.full.return_value = False
read_ahead_cursor = ReadAheadCursor(MOCK_STATEMENT_RESULT,
mock_session, MOCK_TRANSACTION_ID,
MOCK_READ_AHEAD, mock_executor)
read_ahead_cursor._queue = mock_queue
read_ahead_cursor._populate_queue()
next_page_token = MOCK_STATEMENT_RESULT.get('FirstPage').get(
'NextPageToken')
mock_session._fetch_page.assert_called_once_with(
MOCK_TRANSACTION_ID, next_page_token)
mock_queue.put.assert_called_once_with(
MOCK_FETCH_PAGE_RESULT_WITHOUT_TOKEN, timeout=0.05)
def test_populate_queue_result_closed_error(self, mock_session,
mock_logger_debug,
mock_executor):
def close_parent_txn(txn_id, token):
read_ahead_cursor._is_open = False
return MOCK_STATEMENT_RESULT
mock_logger_debug.return_value = None
mock_session.return_value = None
mock_session._fetch_page.side_effect = close_parent_txn
read_ahead_cursor = ReadAheadCursor(MOCK_STATEMENT_RESULT,
mock_session, MOCK_TRANSACTION_ID,
MOCK_READ_AHEAD, mock_executor)
read_ahead_cursor._queue = Queue(1)
read_ahead_cursor._queue.put('value to be removed')
read_ahead_cursor._populate_queue()
self.assertEqual(mock_logger_debug.call_count, 2)
self.assertIsInstance(read_ahead_cursor._queue.get(),
ResultClosedError)
self.assertEqual(read_ahead_cursor._queue.qsize(), 0)
def test_are_there_more_results(self, mock_session, mock_executor):
mock_session.return_value = None
read_ahead_cursor = ReadAheadCursor(MOCK_STATEMENT_RESULT,
mock_session, MOCK_TRANSACTION_ID,
MOCK_READ_AHEAD, mock_executor)
read_ahead_cursor._page = {'NextPageToken': 'token', 'Values': []}
self.assertTrue(read_ahead_cursor._are_there_more_results())
read_ahead_cursor._page = {'NextPageToken': None, 'Values': []}
read_ahead_cursor._queue = Queue()
self.assertFalse(read_ahead_cursor._are_there_more_results())
def test_read_ahead_queue_with_query_stats(self, mock_session,
mock_executor):
mock_statement_result_1 = generate_statement_result(
1, 2, 3, MOCK_TOKEN, True)
mock_statement_result_2 = generate_statement_result(
1, 2, 3, MOCK_TOKEN, False)
mock_statement_result_3 = generate_statement_result(
1, 2, 3, None, False)
def fetch_page(txn_id, token):
statement_results = [
mock_statement_result_2, mock_statement_result_3
]
statement_result = statement_results[fetch_page.page_num]
fetch_page.page_num += 1
return statement_result
fetch_page.page_num = 0
mock_session.return_value = None
mock_session._fetch_page.side_effect = fetch_page
mock_read_ahead = 3
read_ahead_cursor = ReadAheadCursor(mock_statement_result_1,
mock_session, MOCK_TRANSACTION_ID,
mock_read_ahead, mock_executor)
read_ahead_cursor._value_holder_to_ion_value = MagicMock(
name='_value_holder_to_ion_value')
read_ahead_cursor._populate_queue()
# Queue should be populated with the next two pages
self.assertEqual(read_ahead_cursor._queue.qsize(), mock_read_ahead - 1)
# Even if queue is populated with the next two pages, query stats should only total the first page here
assert_query_stats(self, read_ahead_cursor, 1, 2, 3)
read_ahead_cursor._next_page()
# Query stats should only total the first page and second page here
assert_query_stats(self, read_ahead_cursor, 2, 4, 6)
read_ahead_cursor._next_page()
# Query stats should total all three pages
assert_query_stats(self, read_ahead_cursor, 3, 6, 9)
def test_next(self, mock_session, mock_value_holder_to_ion_value,
mock_executor):
mock_session.return_value = None
mock_value_holder_to_ion_value.side_effect = lambda val: val
read_ahead_cursor = ReadAheadCursor(
MOCK_STATEMENT_RESULT_WITHOUT_TOKEN, mock_session,
MOCK_TRANSACTION_ID, MOCK_READ_AHEAD, mock_executor)
count = 0
for value in MOCK_ION_BINARY_VALUES:
self.assertEqual(read_ahead_cursor._index, count)
self.assertEqual(next(read_ahead_cursor), value)
mock_value_holder_to_ion_value.assert_called_with(value)
count += 1
self.assertRaises(StopIteration, next, read_ahead_cursor)
def test_next(self, mock_session, mock_value_holder_to_ion_value,
mock_executor):
mock_session.return_value = None
mock_value_holder_to_ion_value.side_effect = lambda val: val
mock_statement_result_with_none_next_page_token = MOCK_STATEMENT_RESULT.copy(
)
mock_statement_result_with_none_next_page_token.update(
{'NextPageToken': None})
read_ahead_cursor = ReadAheadCursor(
mock_statement_result_with_none_next_page_token, mock_session,
MOCK_TRANSACTION_ID, MOCK_READ_AHEAD, mock_executor)
count = 0
for value in MOCK_ION_BINARY_VALUES:
self.assertEqual(read_ahead_cursor._index, count)
self.assertEqual(next(read_ahead_cursor), value)
mock_value_holder_to_ion_value.assert_called_with(value)
count += 1
self.assertRaises(StopIteration, next, read_ahead_cursor)
def test_ReadAheadCursor_with_executor(self, mock_queue, mock_session,
mock_executor):
mock_session.return_value = None
mock_queue.return_value = mock_queue
read_ahead_cursor = ReadAheadCursor(MOCK_STATEMENT_RESULT,
mock_session, MOCK_TRANSACTION_ID,
MOCK_READ_AHEAD, mock_executor)
self.assertEqual(read_ahead_cursor._page, MOCK_STATEMENT_RESULT)
self.assertEqual(read_ahead_cursor._session, mock_session)
self.assertEqual(read_ahead_cursor._transaction_id,
MOCK_TRANSACTION_ID)
self.assertEqual(read_ahead_cursor._index, 0)
self.assertEqual(read_ahead_cursor._queue, mock_queue)
self.assertEqual(read_ahead_cursor._is_open, True)
mock_queue.assert_called_once_with(MOCK_READ_AHEAD - 1)
mock_executor.submit.assert_called_once_with(
read_ahead_cursor._populate_queue)
def test_ReadAHeadCursor_without_executor(self, mock_thread, mock_queue,
mock_session):
mock_session.return_value = None
mock_thread.return_value = mock_thread
mock_queue.return_value = mock_queue
read_ahead_cursor = ReadAheadCursor(MOCK_STATEMENT_RESULT,
mock_session, MOCK_TRANSACTION_ID,
MOCK_READ_AHEAD, None)
self.assertEqual(read_ahead_cursor._page, MOCK_STATEMENT_RESULT)
self.assertEqual(read_ahead_cursor._session, mock_session)
self.assertEqual(read_ahead_cursor._transaction_id,
MOCK_TRANSACTION_ID)
self.assertEqual(read_ahead_cursor._index, 0)
self.assertEqual(read_ahead_cursor._queue, mock_queue)
self.assertEqual(read_ahead_cursor._is_open, True)
mock_queue.assert_called_once_with(MOCK_READ_AHEAD - 1)
mock_thread.assert_called_once_with(
target=read_ahead_cursor._populate_queue)
mock_thread().setDaemon.assert_called_once_with(True)
mock_thread().start.assert_called_once_with()
def test_iter(self, mock_session, mock_executor):
mock_session.return_value = None
read_ahead_cursor = ReadAheadCursor(MOCK_STATEMENT_RESULT,
mock_session, MOCK_TRANSACTION_ID,
MOCK_READ_AHEAD, mock_executor)
self.assertEqual(iter(read_ahead_cursor), read_ahead_cursor)
def test_value_holder_to_ion_value(self):
ion_value = 'IonValue'
value_holder = {'IonBinary': dumps(ion_value)}
result = ReadAheadCursor._value_holder_to_ion_value(value_holder)
self.assertEqual(result, ion_value)
def test_next_when_closed(self, mock_session, mock_executor):
read_ahead_cursor = ReadAheadCursor(MOCK_STATEMENT_RESULT,
mock_session, MOCK_TRANSACTION_ID,
MOCK_READ_AHEAD, mock_executor)
read_ahead_cursor.close()
self.assertRaises(ResultClosedError, next, read_ahead_cursor)
