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 test_order_by():
query = query_module.QueryOptions(
order_by=[
query_module.PropertyOrder("a"),
query_module.PropertyOrder("b", reverse=True),
]
)
expected_pb = query_pb2.Query(
order=[
query_pb2.PropertyOrder(
property=query_pb2.PropertyReference(name="a"),
direction=query_pb2.PropertyOrder.ASCENDING,
),
query_pb2.PropertyOrder(
property=query_pb2.PropertyReference(name="b"),
direction=query_pb2.PropertyOrder.DESCENDING,
),
]
)
assert _datastore_query._query_to_protobuf(query) == expected_pb
def test_iterate_async_w_limit_and_offset():
def predicate(result):
return result.result % 2 == 0
query = query_module.QueryOptions(offset=1, limit=2)
iterator = _datastore_query._PostFilterQueryIteratorImpl(
query, predicate)
iterator._result_set = MockResultSet([1, 2, 3, 4, 5, 6, 7, 8])
@tasklets.tasklet
def iterate():
results = []
while (yield iterator.has_next_async()):
results.append(iterator.next())
raise tasklets.Return(results)
assert iterate().result() == [4, 6]
with pytest.raises(StopIteration):
iterator.next()
def test_iterate_async_ordered_limit_and_offset():
foo = model.StringProperty("foo")
query = query_module.QueryOptions(
offset=5,
limit=4,
filters=query_module.OR(foo == "this", foo == "that"),
)
iterator = _datastore_query._MultiQueryIteratorImpl(query)
iterator._sortable = True
iterator._result_sets = [
MockResultSet(["a", "c", "e", "g", "i"]),
MockResultSet(["a", "b", "d", "f", "h", "j"]),
]
@tasklets.tasklet
def iterate():
results = []
while (yield iterator.has_next_async()):
results.append(iterator.next())
raise tasklets.Return(results)
assert iterate().result() == ["f", "g", "h", "i"]
def test_ancestor_with_composite_filter():
key = key_module.Key("Foo", 123)
foo = model.StringProperty("foo")
food = model.StringProperty("food")
query = query_module.QueryOptions(
ancestor=key,
filters=query_module.AND(foo == "bar", food == "barn"),
)
query_pb = _datastore_query._query_to_protobuf(query)
filter_pb1 = query_pb2.PropertyFilter(
property=query_pb2.PropertyReference(name="foo"),
op=query_pb2.PropertyFilter.EQUAL,
value=entity_pb2.Value(string_value="bar"),
)
filter_pb2 = query_pb2.PropertyFilter(
property=query_pb2.PropertyReference(name="food"),
op=query_pb2.PropertyFilter.EQUAL,
value=entity_pb2.Value(string_value="barn"),
)
ancestor_pb = query_pb2.PropertyFilter(
property=query_pb2.PropertyReference(name="__key__"),
op=query_pb2.PropertyFilter.HAS_ANCESTOR,
)
ancestor_pb.value.key_value.CopyFrom(key._key.to_protobuf())
expected_pb = query_pb2.Query(
filter=query_pb2.Filter(
composite_filter=query_pb2.CompositeFilter(
op=query_pb2.CompositeFilter.AND,
filters=[
query_pb2.Filter(property_filter=filter_pb1),
query_pb2.Filter(property_filter=filter_pb2),
query_pb2.Filter(property_filter=ancestor_pb),
],
)
)
)
assert query_pb == expected_pb
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_cursor_after_no_cursor():
query = query_module.QueryOptions()
iterator = _datastore_query._PostFilterQueryIteratorImpl(
query, "predicate")
with pytest.raises(exceptions.BadArgumentError):
iterator.cursor_after()
def test_limit():
query = query_module.QueryOptions(limit=20)
expected_pb = query_pb2.Query()
expected_pb.limit.value = 20
assert _datastore_query._query_to_protobuf(query) == expected_pb
def test_offset():
query = query_module.QueryOptions(offset=20)
assert _datastore_query._query_to_protobuf(query) == query_pb2.Query(
offset=20
)
def test_end_cursor():
query = query_module.QueryOptions(
end_cursor=_datastore_query.Cursor(b"abc"))
assert _datastore_query._query_to_protobuf(query) == query_pb2.Query(
end_cursor=b"abc")
def test_kind():
query = query_module.QueryOptions(kind="Foo")
assert _datastore_query._query_to_protobuf(query) == query_pb2.Query(
kind=[query_pb2.KindExpression(name="Foo")]
)
def test_no_args():
query = query_module.QueryOptions()
assert _datastore_query._query_to_protobuf(query) == query_pb2.Query()
def test_iter():
foo = model.StringProperty("foo")
query = query_module.QueryOptions(
filters=query_module.OR(foo == "this", foo == "that"))
iterator = _datastore_query._MultiQueryIteratorImpl(query)
assert iter(iterator) is iterator
def test_constructor():
with pytest.raises(NotImplementedError):
query.QueryOptions()
def test_cursor_after():
query = query_module.QueryOptions()
iterator = _datastore_query._PostFilterQueryIteratorImpl(
query, "predicate")
iterator._cursor_after = "himom"
assert iterator.cursor_after() == "himom"
def test_probably_has_next_delegate():
query = query_module.QueryOptions()
iterator = _datastore_query._PostFilterQueryIteratorImpl(
query, "predicate")
iterator._result_set._next_result = "foo"
assert iterator.probably_has_next() is True
def test_has_next_async_next_loaded():
query = query_module.QueryOptions()
iterator = _datastore_query._PostFilterQueryIteratorImpl(
query, "predicate")
iterator._next_result = "foo"
assert iterator.has_next_async().result()
