def test_constructor_raw():
iterator = _datastore_query._QueryIteratorImpl("foo", raw=True)
assert iterator._query == "foo"
assert iterator._batch is None
assert iterator._index is None
assert iterator._has_next_batch is None
assert iterator._cursor_before is None
assert iterator._cursor_after is None
assert iterator._raw
def test_next_raw():
iterator = _datastore_query._QueryIteratorImpl("foo", raw=True)
iterator.has_next = mock.Mock(return_value=True)
iterator._index = 0
result = mock.Mock(cursor=b"abc")
iterator._batch = [result]
assert iterator.next() is result
assert iterator._index == 1
assert iterator._cursor_after == b"abc"
def test_next_entity():
iterator = _datastore_query._QueryIteratorImpl("foo")
iterator.has_next = mock.Mock(return_value=True)
iterator._index = 1
iterator._cursor_before = b"abc"
result = mock.Mock(cursor=b"bcd")
iterator._batch = [None, result]
assert iterator.next() is result.entity.return_value
assert iterator._index == 2
assert iterator._cursor_after == b"bcd"
def test_has_next_async_not_started():
iterator = _datastore_query._QueryIteratorImpl("foo")
def dummy_next_batch():
iterator._index = 0
iterator._batch = ["a", "b", "c"]
return utils.future_result(None)
iterator._next_batch = dummy_next_batch
assert iterator.has_next_async().result()
def test_next_done():
iterator = _datastore_query._QueryIteratorImpl("foo")
iterator.has_next = mock.Mock(return_value=False)
iterator._cursor_before = b"abc"
iterator._cursor_after = b"bcd"
with pytest.raises(StopIteration):
iterator.next()
with pytest.raises(exceptions.BadArgumentError):
iterator.cursor_before()
assert iterator.cursor_after() == b"bcd"
def test_has_next_async_next_batch_finished():
iterator = _datastore_query._QueryIteratorImpl("foo")
iterator._index = 3
iterator._batch = ["a", "b", "c"]
iterator._has_next_batch = True
def dummy_next_batch():
iterator._index = 3
iterator._batch = ["d", "e", "f"]
return utils.future_result(None)
iterator._next_batch = dummy_next_batch
assert not iterator.has_next_async().result()
def test_has_next_async_next_batch_is_empty():
iterator = _datastore_query._QueryIteratorImpl("foo")
iterator._index = 3
iterator._batch = ["a", "b", "c"]
iterator._has_next_batch = True
batches = [[], ["d", "e", "f"]]
def dummy_next_batch():
iterator._index = 0
iterator._batch = batches.pop(0)
return utils.future_result(None)
iterator._next_batch = dummy_next_batch
assert iterator.has_next_async().result()
def test__next_batch(_datastore_run_query):
entity_results = [
mock.Mock(entity="entity1", cursor=b"a"),
mock.Mock(entity="entity2", cursor=b"b"),
mock.Mock(entity="entity3", cursor=b"c"),
]
_datastore_run_query.return_value = utils.future_result(
mock.Mock(batch=mock.Mock(
entity_result_type=query_pb2.EntityResult.FULL,
entity_results=entity_results,
end_cursor=b"abc",
more_results=query_pb2.QueryResultBatch.NO_MORE_RESULTS,
)))
query = query_module.QueryOptions()
iterator = _datastore_query._QueryIteratorImpl(query)
assert iterator._next_batch().result() is None
assert iterator._index == 0
assert len(iterator._batch) == 3
assert iterator._batch[0].result_pb.entity == "entity1"
assert iterator._batch[0].result_type == query_pb2.EntityResult.FULL
assert iterator._batch[0].order_by is None
assert not iterator._has_next_batch
def test__peek():
iterator = _datastore_query._QueryIteratorImpl("foo")
iterator._index = 1
iterator._batch = ["a", "b", "c"]
assert iterator._peek() == "b"
def test_probably_has_next_finished():
iterator = _datastore_query._QueryIteratorImpl("foo")
iterator._batch = ["a", "b", "c"]
iterator._index = 3
assert not iterator.probably_has_next()
def test_cursor_after_no_cursor():
iterator = _datastore_query._QueryIteratorImpl("foo")
with pytest.raises(exceptions.BadArgumentError):
iterator.cursor_after()
def test_has_next_async_finished():
iterator = _datastore_query._QueryIteratorImpl("foo")
iterator._index = 3
iterator._batch = ["a", "b", "c"]
assert not iterator.has_next_async().result()
def test_has_next():
iterator = _datastore_query._QueryIteratorImpl("foo")
iterator.has_next_async = mock.Mock(
return_value=utils.future_result("bar")
)
assert iterator.has_next() == "bar"
def test___iter__():
iterator = _datastore_query._QueryIteratorImpl("foo")
assert iter(iterator) is iterator
def test__peek_key_error():
iterator = _datastore_query._QueryIteratorImpl("foo")
with pytest.raises(KeyError):
iterator._peek()
def test_cursor_after():
iterator = _datastore_query._QueryIteratorImpl("foo")
iterator._cursor_after = "foo"
assert iterator.cursor_after() == "foo"
def test_probably_has_next_not_started():
iterator = _datastore_query._QueryIteratorImpl("foo")
assert iterator.probably_has_next()
def test_index_list():
iterator = _datastore_query._QueryIteratorImpl("foo")
with pytest.raises(NotImplementedError):
iterator.index_list()
def test_probably_has_next_more_batches():
iterator = _datastore_query._QueryIteratorImpl("foo")
iterator._batch = "foo"
iterator._has_next_batch = True
assert iterator.probably_has_next()
