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