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