def test__next_batch_has_more_w_offset_and_limit(_datastore_run_query):
"""Regression test for Issue #236
https://github.com/googleapis/python-ndb/issues/236
"""
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.NOT_FINISHED,
)
)
)
query = query_module.QueryOptions(offset=5, limit=5)
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 iterator._has_next_batch
assert iterator._query.start_cursor.cursor == b"abc"
assert iterator._query.offset is None
assert iterator._query.limit == 2
def test_has_next():
query = query_module.QueryOptions()
iterator = _datastore_query._PostFilterQueryIteratorImpl(
query, "predicate")
iterator.has_next_async = mock.Mock(
return_value=utils.future_result("bar"))
assert iterator.has_next() == "bar"
def test__next_batch_has_more(_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.NOT_FINISHED,
)
)
)
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 iterator._has_next_batch
assert iterator._query.start_cursor.cursor == b"abc"
def test_has_next():
foo = model.StringProperty("foo")
query = query_module.QueryOptions(
filters=query_module.OR(foo == "this", foo == "that"))
iterator = _datastore_query._MultiQueryIteratorImpl(query)
iterator.has_next_async = mock.Mock(
return_value=utils.future_result("bar"))
assert iterator.has_next() == "bar"
def test_allocate_ids_callback(_datastore_allocate_ids):
options = _options.Options()
batch = _api._AllocateIdsBatch(options)
batch.futures = futures = [tasklets.Future(), tasklets.Future()]
rpc = utils.future_result(
mock.Mock(keys=["key1", "key2"], spec=("key", )))
batch.allocate_ids_callback(rpc)
results = [future.result() for future in futures]
assert results == ["key1", "key2"]
def test_it(_datastore_api):
query = query_module.QueryOptions(project="testing", namespace="")
query_pb = _datastore_query._query_to_protobuf(query)
_datastore_api.make_call.return_value = utils.future_result("foo")
read_options = datastore_pb2.ReadOptions()
request = datastore_pb2.RunQueryRequest(
project_id="testing",
partition_id=entity_pb2.PartitionId(project_id="testing",
namespace_id=""),
query=query_pb,
read_options=read_options,
)
_datastore_api.get_read_options.return_value = read_options
assert _datastore_query._datastore_run_query(query).result() == "foo"
_datastore_api.make_call.assert_called_once_with("RunQuery",
request,
timeout=None)
_datastore_api.get_read_options.assert_called_once_with(
query, default_read_consistency=_datastore_api.EVENTUAL)
def has_next_async(self):
return utils.future_result(self.index < self.len)
def dummy_next_batch():
iterator._index = 3
iterator._batch = ["d", "e", "f"]
return utils.future_result(None)
def dummy_next_batch():
iterator._index = 0
iterator._batch = ["a", "b", "c"]
return utils.future_result(None)
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 dummy_next_batch():
iterator._index = 0
iterator._batch = batches.pop(0)
return utils.future_result(None)
