def test__items_iter(self):
# Items to be returned.
item1 = 17
item2 = 100
item3 = 211
# Make pages from mock responses
parent = mock.sentinel.parent
page1 = page_iterator.Page(parent, (item1, item2),
page_iterator._item_to_value_identity)
page2 = page_iterator.Page(parent, (item3, ),
page_iterator._item_to_value_identity)
iterator = PageIteratorImpl(None, None)
iterator._next_page = mock.Mock(side_effect=[page1, page2, None])
items_iter = iterator._items_iter()
assert isinstance(items_iter, types.GeneratorType)
# Consume items and check the state of the iterator.
assert iterator.num_results == 0
assert six.next(items_iter) == item1
assert iterator.num_results == 1
assert six.next(items_iter) == item2
assert iterator.num_results == 2
assert six.next(items_iter) == item3
assert iterator.num_results == 3
with pytest.raises(StopIteration):
six.next(items_iter)
def _next_page(self):
"""Get the next page in the iterator.
:rtype: :class:`~google.cloud.iterator.Page`
:returns: The next page in the iterator (or :data:`None` if
there are no pages left).
"""
if not self._more_results:
return None
query_pb = self._build_protobuf()
transaction = self.client.current_transaction
if transaction is None:
read_options = _datastore_pb2.ReadOptions()
else:
read_options = _datastore_pb2.ReadOptions(
transaction=transaction.id)
partition_id = _entity_pb2.PartitionId(
project_id=self._query.project, namespace_id=self._query.namespace)
response_pb = self.client._datastore_api.run_query(
self._query.project,
partition_id,
read_options,
query=query_pb,
)
entity_pbs = self._process_query_results(response_pb)
return page_iterator.Page(self, entity_pbs, self._item_to_value)
def test_iterator_calls_parent_item_to_value(self):
parent = mock.sentinel.parent
item_to_value = mock.Mock(side_effect=lambda iterator, value: value,
spec=["__call__"])
page = page_iterator.Page(parent, (10, 11, 12), item_to_value)
page._remaining = 100
assert item_to_value.call_count == 0
assert page.remaining == 100
assert six.next(page) == 10
assert item_to_value.call_count == 1
item_to_value.assert_called_with(parent, 10)
assert page.remaining == 99
assert six.next(page) == 11
assert item_to_value.call_count == 2
item_to_value.assert_called_with(parent, 11)
assert page.remaining == 98
assert six.next(page) == 12
assert item_to_value.call_count == 3
item_to_value.assert_called_with(parent, 12)
assert page.remaining == 97
def test_next(self):
iterator = PageIteratorImpl(None, None)
page_1 = page_iterator.Page(iterator, ("item 1.1", "item 1.2"),
page_iterator._item_to_value_identity)
page_2 = page_iterator.Page(iterator, ("item 2.1", ),
page_iterator._item_to_value_identity)
iterator._next_page = mock.Mock(side_effect=[page_1, page_2, None])
result = next(iterator)
assert result == "item 1.1"
result = next(iterator)
assert result == "item 1.2"
result = next(iterator)
assert result == "item 2.1"
with pytest.raises(StopIteration):
next(iterator)
def _next_page(self):
"""Get the next page in the iterator.
:rtype: :class:`~google.cloud.iterator.Page`
:returns: The next page in the iterator (or :data:`None` if
there are no pages left).
"""
if not self._more_results:
return None
query_pb = self._build_protobuf()
transaction = self.client.current_transaction
if transaction is None:
transaction_id = None
else:
transaction_id = transaction.id
read_options = helpers.get_read_options(self._eventual, transaction_id)
partition_id = entity_pb2.PartitionId(
project_id=self._query.project, namespace_id=self._query.namespace)
kwargs = {}
if self._retry is not None:
kwargs["retry"] = self._retry
if self._timeout is not None:
kwargs["timeout"] = self._timeout
response_pb = self.client._datastore_api.run_query(
request={
"project_id": self._query.project,
"partition_id": partition_id,
"read_options": read_options,
"query": query_pb,
},
**kwargs,
)
while (response_pb.batch.more_results == _NOT_FINISHED
and response_pb.batch.skipped_results < query_pb.offset):
# We haven't finished processing. A likely reason is we haven't
# skipped all of the results yet. Don't return any results.
# Instead, rerun query, adjusting offsets. Datastore doesn't process
# more than 1000 skipped results in a query.
query_pb.start_cursor = response_pb.batch.skipped_cursor
query_pb.offset -= response_pb.batch.skipped_results
response_pb = self.client._datastore_api.run_query(
request={
"project_id": self._query.project,
"partition_id": partition_id,
"read_options": read_options,
"query": query_pb,
},
**kwargs,
)
entity_pbs = self._process_query_results(response_pb)
return page_iterator.Page(self, entity_pbs, self.item_to_value)
def test_list_buckets_page_empty_response(self):
from google.api_core import page_iterator
project = "PROJECT"
credentials = _make_credentials()
client = self._make_one(project=project, credentials=credentials)
iterator = client.list_buckets()
page = page_iterator.Page(iterator, (), None)
iterator._page = page
self.assertEqual(list(page), [])
def test_constructor(self):
parent = mock.sentinel.parent
item_to_value = mock.sentinel.item_to_value
page = page_iterator.Page(parent, (1, 2, 3), item_to_value)
assert page.num_items == 3
assert page.remaining == 3
assert page._parent is parent
assert page._item_to_value is item_to_value
def test_page_empty_response(self):
from google.api_core import page_iterator
credentials = _make_credentials()
client = self._make_one(credentials=credentials)
iterator = client.list_projects()
page = page_iterator.Page(iterator, (), None)
iterator._page = page
self.assertEqual(page.num_items, 0)
self.assertEqual(page.remaining, 0)
self.assertEqual(list(page), [])
def test_raw_page(self):
parent = mock.sentinel.parent
item_to_value = mock.sentinel.item_to_value
raw_page = mock.sentinel.raw_page
page = page_iterator.Page(parent, (1, 2, 3), item_to_value, raw_page=raw_page)
assert page.raw_page is raw_page
with pytest.raises(AttributeError):
page.raw_page = None
def test__page_iter_increment(self):
iterator = PageIteratorImpl(None, None)
page = page_iterator.Page(iterator, ("item", ),
page_iterator._item_to_value_identity)
iterator._next_page = mock.Mock(side_effect=[page, None])
assert iterator.num_results == 0
page_iter = iterator._page_iter(increment=True)
next(page_iter)
assert iterator.num_results == 1
def test___iter__(self):
page = page_iterator.Page(None, (), None, None)
assert iter(page) is page
