def test_nested_entity_no_key(self):
from google.cloud.datastore_v1.types import entity as entity_pb2
from google.cloud.datastore.helpers import _new_value_pb
PROJECT = "FOO"
KIND = "KIND"
INSIDE_NAME = "IFOO"
OUTSIDE_NAME = "OBAR"
INSIDE_VALUE = 1337
entity_inside = entity_pb2.Entity()
inside_val_pb = _new_value_pb(entity_inside, INSIDE_NAME)
inside_val_pb.integer_value = INSIDE_VALUE
entity_pb = entity_pb2.Entity()
entity_pb.key.partition_id.project_id = PROJECT
element = entity_pb._pb.key.path.add()
element.kind = KIND
outside_val_pb = _new_value_pb(entity_pb, OUTSIDE_NAME)
outside_val_pb.entity_value.CopyFrom(entity_inside._pb)
entity = self._call_fut(entity_pb._pb)
self.assertEqual(entity.key.project, PROJECT)
self.assertEqual(entity.key.flat_path, (KIND, ))
self.assertEqual(len(entity), 1)
inside_entity = entity[OUTSIDE_NAME]
self.assertIsNone(inside_entity.key)
self.assertEqual(len(inside_entity), 1)
self.assertEqual(inside_entity[INSIDE_NAME], INSIDE_VALUE)
def test_dict_to_entity_recursive(self):
from google.cloud.datastore_v1.types import entity as entity_pb2
from google.cloud.datastore.entity import Entity
entity = Entity()
entity["a"] = {"b": {"c": {"d": 1.25}, "e": True}, "f": 10}
entity_pb = self._call_fut(entity)
b_entity_pb = entity_pb2.Entity(
properties={
"c":
entity_pb2.Value(entity_value=entity_pb2.Entity(
properties={"d": entity_pb2.Value(double_value=1.25)})),
"e":
entity_pb2.Value(boolean_value=True),
})
expected_pb = entity_pb2.Entity(
properties={
"a":
entity_pb2.Value(entity_value=entity_pb2.Entity(
properties={
"b": entity_pb2.Value(entity_value=b_entity_pb),
"f": entity_pb2.Value(integer_value=10),
}))
})
self.assertEqual(entity_pb, expected_pb)
def test_inverts_to_protobuf(self):
from google.cloud.datastore_v1.types import entity as entity_pb2
from google.cloud.datastore.helpers import _new_value_pb
from google.cloud.datastore.helpers import entity_from_protobuf
original_pb = entity_pb2.Entity()
# Add a key.
original_pb.key.partition_id.project_id = project = "PROJECT"
elem1 = original_pb._pb.key.path.add()
elem1.kind = "Family"
elem1.id = 1234
elem2 = original_pb._pb.key.path.add()
elem2.kind = "King"
elem2.name = "Spades"
# Add an integer property.
val_pb1 = _new_value_pb(original_pb, "foo")
val_pb1.integer_value = 1337
val_pb1.exclude_from_indexes = True
# Add a string property.
val_pb2 = _new_value_pb(original_pb, "bar")
val_pb2.string_value = u"hello"
# Add a nested (entity) property.
val_pb3 = _new_value_pb(original_pb, "entity-baz")
sub_pb = entity_pb2.Entity()
sub_val_pb1 = _new_value_pb(sub_pb, "x")
sub_val_pb1.double_value = 3.14
sub_val_pb2 = _new_value_pb(sub_pb, "y")
sub_val_pb2.double_value = 2.718281828
val_pb3.meaning = 9
val_pb3.entity_value.CopyFrom(sub_pb._pb)
# Add a list property.
val_pb4 = _new_value_pb(original_pb, "list-quux")
array_val1 = val_pb4.array_value.values.add()
array_val1.exclude_from_indexes = False
array_val1.meaning = meaning = 22
array_val1.blob_value = b"\xe2\x98\x83"
array_val2 = val_pb4.array_value.values.add()
array_val2.exclude_from_indexes = False
array_val2.meaning = meaning
array_val2.blob_value = b"\xe2\x98\x85"
# Convert to the user-space Entity.
entity = entity_from_protobuf(original_pb)
# Convert the user-space Entity back to a protobuf.
new_pb = self._call_fut(entity)
# NOTE: entity_to_protobuf() strips the project so we "cheat".
new_pb.key.partition_id.project_id = project
self._compare_entity_proto(original_pb, new_pb)
def test_dict_to_entity(self):
from google.cloud.datastore_v1.types import entity as entity_pb2
from google.cloud.datastore.entity import Entity
entity = Entity()
entity["a"] = {"b": u"c"}
entity_pb = self._call_fut(entity)
expected_pb = entity_pb2.Entity(
properties={
"a":
entity_pb2.Value(entity_value=entity_pb2.Entity(
properties={"b": entity_pb2.Value(string_value="c")}))
})
self.assertEqual(entity_pb, expected_pb)
def test_variable_meanings(self):
from google.cloud.datastore_v1.types import entity as entity_pb2
from google.cloud.datastore.entity import Entity
from google.cloud.datastore.helpers import _new_value_pb
entity = Entity()
name = "quux"
entity[name] = values = [1, 20, 300]
meaning = 9
entity._meanings[name] = ([None, meaning, None], values)
entity_pb = self._call_fut(entity)
# Construct the expected protobuf.
expected_pb = entity_pb2.Entity()
value_pb = _new_value_pb(expected_pb, name)
value0 = value_pb.array_value.values.add()
value0.integer_value = values[0]
# The only array entry with a meaning is the middle one.
value1 = value_pb.array_value.values.add()
value1.integer_value = values[1]
value1.meaning = meaning
value2 = value_pb.array_value.values.add()
value2.integer_value = values[2]
self._compare_entity_proto(entity_pb, expected_pb)
def _deserialize_entity(pb_bytes):
"""Converts a serialized proto string to an Entity object."""
if pb_bytes is None:
return None
entity = entity_pb2.Entity()
entity._pb.ParseFromString(pb_bytes)
return helpers.entity_from_protobuf(entity._pb)
def test_empty(self):
from google.cloud.datastore_v1.types import entity as entity_pb2
from google.cloud.datastore.entity import Entity
entity = Entity()
entity_pb = self._call_fut(entity)
self._compare_entity_proto(entity_pb, entity_pb2.Entity())
def entity_to_protobuf(entity):
"""Converts an entity into a protobuf.
:type entity: :class:`google.cloud.datastore.entity.Entity`
:param entity: The entity to be turned into a protobuf.
:rtype: :class:`.entity_pb2.Entity`
:returns: The protobuf representing the entity.
"""
entity_pb = entity_pb2.Entity()
if entity.key is not None:
key_pb = entity.key.to_protobuf()
entity_pb._pb.key.CopyFrom(key_pb._pb)
for name, value in entity.items():
value_is_list = isinstance(value, list)
value_pb = _new_value_pb(entity_pb, name)
# Set the appropriate value.
_set_protobuf_value(value_pb, value)
# Add index information to protobuf.
if name in entity.exclude_from_indexes:
if not value_is_list:
value_pb.exclude_from_indexes = True
for sub_value in value_pb.array_value.values:
sub_value.exclude_from_indexes = True
# Add meaning information to protobuf.
_set_pb_meaning_from_entity(
entity, name, value, value_pb, is_list=value_is_list
)
return entity_pb
def test_pb2_entity_no_key(self):
from google.cloud.datastore_v1.types import entity as entity_pb2
entity_pb = entity_pb2.Entity()
entity = self._call_fut(entity_pb)
self.assertIsNone(entity.key)
self.assertEqual(dict(entity), {})
def _make_entity(kind, id_, project):
from google.cloud.datastore_v1.types import entity as entity_pb2
key = entity_pb2.Key()
key.partition_id.project_id = project
elem = key.path._pb.add()
elem.kind = kind
elem.id = id_
return entity_pb2.Entity(key=key)
def test_it(self):
from google.cloud.datastore_v1.types import entity as entity_pb2
entity_pb = entity_pb2.Entity()
name = "foo"
result = self._call_fut(entity_pb, name)
self.assertIsInstance(result, type(entity_pb2.Value()._pb))
self.assertEqual(len(entity_pb._pb.properties), 1)
self.assertEqual(entity_pb._pb.properties[name], result)
def test_with_empty_list(self):
from google.cloud.datastore_v1.types import entity as entity_pb2
from google.cloud.datastore.entity import Entity
entity = Entity()
entity["foo"] = []
entity_pb = self._call_fut(entity)
expected_pb = entity_pb2.Entity()
prop = expected_pb._pb.properties.get_or_create("foo")
prop.array_value.CopyFrom(entity_pb2.ArrayValue(values=[])._pb)
self._compare_entity_proto(entity_pb, expected_pb)
def test_entity_with_meaning(self):
from google.cloud.datastore_v1.types import entity as entity_pb2
from google.cloud.datastore.helpers import _new_value_pb
entity_pb = entity_pb2.Entity()
name = "hello"
value_pb = _new_value_pb(entity_pb, name)
value_pb.meaning = meaning = 9
value_pb.string_value = val = u"something"
entity = self._call_fut(entity_pb)
self.assertIsNone(entity.key)
self.assertEqual(dict(entity), {name: val})
self.assertEqual(entity._meanings, {name: (meaning, val)})
def test_it(self):
from google.cloud.datastore_v1.types import entity as entity_pb2
from google.cloud.datastore.helpers import _new_value_pb
_PROJECT = "PROJECT"
_KIND = "KIND"
_ID = 1234
entity_pb = entity_pb2.Entity()
entity_pb.key.partition_id.project_id = _PROJECT
entity_pb._pb.key.path.add(kind=_KIND, id=_ID)
value_pb = _new_value_pb(entity_pb, "foo")
value_pb.string_value = "Foo"
unindexed_val_pb = _new_value_pb(entity_pb, "bar")
unindexed_val_pb.integer_value = 10
unindexed_val_pb.exclude_from_indexes = True
array_val_pb1 = _new_value_pb(entity_pb, "baz")
array_pb1 = array_val_pb1.array_value.values
unindexed_array_val_pb = array_pb1.add()
unindexed_array_val_pb.integer_value = 11
unindexed_array_val_pb.exclude_from_indexes = True
array_val_pb2 = _new_value_pb(entity_pb, "qux")
array_pb2 = array_val_pb2.array_value.values
indexed_array_val_pb = array_pb2.add()
indexed_array_val_pb.integer_value = 12
entity = self._call_fut(entity_pb._pb)
self.assertEqual(entity.kind, _KIND)
self.assertEqual(entity.exclude_from_indexes, frozenset(["bar",
"baz"]))
entity_props = dict(entity)
self.assertEqual(entity_props, {
"foo": "Foo",
"bar": 10,
"baz": [11],
"qux": [12]
})
# Also check the key.
key = entity.key
self.assertEqual(key.project, _PROJECT)
self.assertIsNone(key.namespace)
self.assertEqual(key.kind, _KIND)
self.assertEqual(key.id, _ID)
def test_it(self):
import types
from google.cloud.datastore_v1.types import entity as entity_pb2
from google.cloud.datastore.helpers import _new_value_pb
entity_pb = entity_pb2.Entity()
name1 = "foo"
name2 = "bar"
val_pb1 = _new_value_pb(entity_pb, name1)
val_pb2 = _new_value_pb(entity_pb, name2)
result = self._call_fut(entity_pb)
self.assertIsInstance(result, types.GeneratorType)
self.assertEqual(sorted(result),
sorted([(name1, val_pb1), (name2, val_pb2)]))
def test_run_query_w_namespace_nonempty_result(self):
from google.cloud.datastore_v1.types import datastore as datastore_pb2
from google.cloud.datastore_v1.types import entity as entity_pb2
from google.cloud.datastore_v1.types import query as query_pb2
project = "PROJECT"
kind = "Kind"
namespace = "NS"
query_pb = self._make_query_pb(kind)
partition_id = entity_pb2.PartitionId(project_id=project,
namespace_id=namespace)
read_options = datastore_pb2.ReadOptions()
rsp_pb = datastore_pb2.RunQueryResponse(
batch=query_pb2.QueryResultBatch(
entity_result_type=query_pb2.EntityResult.ResultType.FULL,
entity_results=[
query_pb2.EntityResult(entity=entity_pb2.Entity())
],
more_results=query_pb2.QueryResultBatch.MoreResultsType.
NO_MORE_RESULTS,
))
# Create mock HTTP and client with response.
http = _make_requests_session(
[_make_response(content=rsp_pb._pb.SerializeToString())])
client_info = _make_client_info()
client = mock.Mock(
_http=http,
_base_url="test.invalid",
_client_info=client_info,
spec=["_http", "_base_url", "_client_info"],
)
# Make request.
ds_api = self._make_one(client)
response = ds_api.run_query(project,
partition_id,
read_options,
query=query_pb)
# Check the result and verify the callers.
self.assertEqual(response, rsp_pb._pb)
uri = _build_expected_url(client._base_url, project, "runQuery")
request = _verify_protobuf_call(http, uri,
datastore_pb2.RunQueryRequest())
self.assertEqual(request.partition_id, partition_id._pb)
self.assertEqual(request.query, query_pb._pb)
def test_meaning_with_change(self):
from google.cloud.datastore_v1.types import entity as entity_pb2
from google.cloud.datastore.entity import Entity
from google.cloud.datastore.helpers import _new_value_pb
entity = Entity()
name = "foo"
entity[name] = value = 42
entity._meanings[name] = (9, 1337)
entity_pb = self._call_fut(entity)
expected_pb = entity_pb2.Entity()
value_pb = _new_value_pb(expected_pb, name)
value_pb.integer_value = value
# NOTE: No meaning is used since the value differs from the
# value stored.
self._compare_entity_proto(entity_pb, expected_pb)
def test_index_mismatch_ignores_empty_list(self):
from google.cloud.datastore_v1.types import entity as entity_pb2
_PROJECT = "PROJECT"
_KIND = "KIND"
_ID = 1234
array_val_pb = entity_pb2.Value(array_value=entity_pb2.ArrayValue(
values=[]))
entity_pb = entity_pb2.Entity(properties={"baz": array_val_pb})
entity_pb.key.partition_id.project_id = _PROJECT
entity_pb.key._pb.path.add(kind=_KIND, id=_ID)
entity = self._call_fut(entity_pb._pb)
entity_dict = dict(entity)
self.assertEqual(entity_dict["baz"], [])
def test_key_only(self):
from google.cloud.datastore_v1.types import entity as entity_pb2
from google.cloud.datastore.entity import Entity
from google.cloud.datastore.key import Key
kind, name = "PATH", "NAME"
project = "PROJECT"
key = Key(kind, name, project=project)
entity = Entity(key=key)
entity_pb = self._call_fut(entity)
expected_pb = entity_pb2.Entity()
expected_pb.key.partition_id.project_id = project
path_elt = expected_pb._pb.key.path.add()
path_elt.kind = kind
path_elt.name = name
self._compare_entity_proto(entity_pb, expected_pb)
def test_simple_fields(self):
from google.cloud.datastore_v1.types import entity as entity_pb2
from google.cloud.datastore.entity import Entity
from google.cloud.datastore.helpers import _new_value_pb
entity = Entity()
name1 = "foo"
entity[name1] = value1 = 42
name2 = "bar"
entity[name2] = value2 = u"some-string"
entity_pb = self._call_fut(entity)
expected_pb = entity_pb2.Entity()
val_pb1 = _new_value_pb(expected_pb, name1)
val_pb1.integer_value = value1
val_pb2 = _new_value_pb(expected_pb, name2)
val_pb2.string_value = value2
self._compare_entity_proto(entity_pb, expected_pb)
def test_lookup_single_key_nonempty_response(self):
from google.cloud.datastore_v1.types import datastore as datastore_pb2
from google.cloud.datastore_v1.types import entity as entity_pb2
project = "PROJECT"
key_pb = _make_key_pb(project)
rsp_pb = datastore_pb2.LookupResponse()
entity = entity_pb2.Entity()
entity.key._pb.CopyFrom(key_pb._pb)
rsp_pb._pb.found.add(entity=entity._pb)
read_options = datastore_pb2.ReadOptions()
# Create mock HTTP and client with response.
http = _make_requests_session(
[_make_response(content=rsp_pb._pb.SerializeToString())])
client_info = _make_client_info()
client = mock.Mock(
_http=http,
_base_url="test.invalid",
_client_info=client_info,
spec=["_http", "_base_url", "_client_info"],
)
# Make request.
ds_api = self._make_one(client)
response = ds_api.lookup(project, [key_pb], read_options=read_options)
# Check the result and verify the callers.
self.assertEqual(response, rsp_pb._pb)
uri = _build_expected_url(client._base_url, project, "lookup")
self.assertEqual(len(response.found), 1)
self.assertEqual(len(response.missing), 0)
self.assertEqual(len(response.deferred), 0)
found = response.found[0].entity
self.assertEqual(found.key.path[0].kind, "Kind")
self.assertEqual(found.key.path[0].id, 1234)
request = _verify_protobuf_call(http, uri,
datastore_pb2.LookupRequest())
self.assertEqual(list(request.keys), [key_pb._pb])
self.assertEqual(request.read_options, read_options._pb)
def test_mismatched_value_indexed(self):
from google.cloud.datastore_v1.types import entity as entity_pb2
from google.cloud.datastore.helpers import _new_value_pb
_PROJECT = "PROJECT"
_KIND = "KIND"
_ID = 1234
entity_pb = entity_pb2.Entity()
entity_pb.key.partition_id.project_id = _PROJECT
entity_pb._pb.key.path.add(kind=_KIND, id=_ID)
array_val_pb = _new_value_pb(entity_pb, "baz")
array_pb = array_val_pb.array_value.values
unindexed_value_pb1 = array_pb.add()
unindexed_value_pb1.integer_value = 10
unindexed_value_pb1.exclude_from_indexes = True
unindexed_value_pb2 = array_pb.add()
unindexed_value_pb2.integer_value = 11
with self.assertRaises(ValueError):
self._call_fut(entity_pb._pb)
def entity_from_protobuf(pb):
"""Factory method for creating an entity based on a protobuf.
The protobuf should be one returned from the Cloud Datastore
Protobuf API.
:type pb: :class:`.entity_pb2.Entity`
:param pb: The Protobuf representing the entity.
:rtype: :class:`google.cloud.datastore.entity.Entity`
:returns: The entity derived from the protobuf.
"""
if not getattr(pb, "_pb", False):
# Coerce raw pb type into proto-plus pythonic type.
proto_pb = entity_pb2.Entity(pb)
pb = pb
else:
proto_pb = pb
pb = pb._pb
key = None
if "key" in proto_pb: # Message field (Key)
key = key_from_protobuf(proto_pb.key)
entity_props = {}
entity_meanings = {}
exclude_from_indexes = []
for prop_name, value_pb in _property_tuples(proto_pb):
value = _get_value_from_value_pb(value_pb)
entity_props[prop_name] = value
# Check if the property has an associated meaning.
is_list = isinstance(value, list)
meaning = _get_meaning(value_pb, is_list=is_list)
if meaning is not None:
entity_meanings[prop_name] = (meaning, value)
# Check if ``value_pb`` was excluded from index. Lists need to be
# special-cased and we require all ``exclude_from_indexes`` values
# in a list agree.
if is_list and len(value) > 0:
exclude_values = set(
value_pb.exclude_from_indexes
for value_pb in value_pb.array_value.values
)
if len(exclude_values) != 1:
raise ValueError(
"For an array_value, subvalues must either "
"all be indexed or all excluded from "
"indexes."
)
if exclude_values.pop():
exclude_from_indexes.append(prop_name)
else:
if value_pb.exclude_from_indexes:
exclude_from_indexes.append(prop_name)
entity = Entity(key=key, exclude_from_indexes=exclude_from_indexes)
entity.update(entity_props)
entity._meanings.update(entity_meanings)
return entity
def _lookup_single_helper(
self,
read_consistency=None,
transaction=None,
empty=True,
retry=None,
timeout=None,
):
from google.cloud.datastore_v1.types import datastore as datastore_pb2
from google.cloud.datastore_v1.types import entity as entity_pb2
project = "PROJECT"
key_pb = _make_key_pb(project)
options_kw = {}
if read_consistency is not None:
options_kw["read_consistency"] = read_consistency
if transaction is not None:
options_kw["transaction"] = transaction
read_options = datastore_pb2.ReadOptions(**options_kw)
rsp_pb = datastore_pb2.LookupResponse()
if not empty:
entity = entity_pb2.Entity()
entity.key._pb.CopyFrom(key_pb._pb)
rsp_pb._pb.found.add(entity=entity._pb)
http = _make_requests_session(
[_make_response(content=rsp_pb._pb.SerializeToString())])
client_info = _make_client_info()
client = mock.Mock(
_http=http,
_base_url="test.invalid",
_client_info=client_info,
spec=["_http", "_base_url", "_client_info"],
)
ds_api = self._make_one(client)
request = {
"project_id": project,
"keys": [key_pb],
"read_options": read_options,
}
kwargs = _make_retry_timeout_kwargs(retry, timeout, http)
response = ds_api.lookup(request=request, **kwargs)
self.assertEqual(response, rsp_pb._pb)
if empty:
self.assertEqual(len(response.found), 0)
else:
self.assertEqual(len(response.found), 1)
self.assertEqual(len(response.missing), 0)
self.assertEqual(len(response.deferred), 0)
uri = _build_expected_url(client._base_url, project, "lookup")
request = _verify_protobuf_call(
http,
uri,
datastore_pb2.LookupRequest(),
retry=retry,
timeout=timeout,
)
if retry is not None:
retry.assert_called_once_with(http.request)
self.assertEqual(list(request.keys), [key_pb._pb])
self.assertEqual(request.read_options, read_options._pb)
def _lookup_multiple_helper(
self,
found=0,
missing=0,
deferred=0,
retry=None,
timeout=None,
):
from google.cloud.datastore_v1.types import datastore as datastore_pb2
from google.cloud.datastore_v1.types import entity as entity_pb2
project = "PROJECT"
key_pb1 = _make_key_pb(project)
key_pb2 = _make_key_pb(project, id_=2345)
keys = [key_pb1, key_pb2]
read_options = datastore_pb2.ReadOptions()
rsp_pb = datastore_pb2.LookupResponse()
found_keys = []
for i_found in range(found):
key = keys[i_found]
found_keys.append(key._pb)
entity = entity_pb2.Entity()
entity.key._pb.CopyFrom(key._pb)
rsp_pb._pb.found.add(entity=entity._pb)
missing_keys = []
for i_missing in range(missing):
key = keys[i_missing]
missing_keys.append(key._pb)
entity = entity_pb2.Entity()
entity.key._pb.CopyFrom(key._pb)
rsp_pb._pb.missing.add(entity=entity._pb)
deferred_keys = []
for i_deferred in range(deferred):
key = keys[i_deferred]
deferred_keys.append(key._pb)
rsp_pb._pb.deferred.append(key._pb)
http = _make_requests_session(
[_make_response(content=rsp_pb._pb.SerializeToString())])
client_info = _make_client_info()
client = mock.Mock(
_http=http,
_base_url="test.invalid",
_client_info=client_info,
spec=["_http", "_base_url", "_client_info"],
)
ds_api = self._make_one(client)
request = {
"project_id": project,
"keys": keys,
"read_options": read_options,
}
kwargs = _make_retry_timeout_kwargs(retry, timeout, http)
response = ds_api.lookup(request=request, **kwargs)
self.assertEqual(response, rsp_pb._pb)
self.assertEqual([found.entity.key for found in response.found],
found_keys)
self.assertEqual([missing.entity.key for missing in response.missing],
missing_keys)
self.assertEqual(list(response.deferred), deferred_keys)
uri = _build_expected_url(client._base_url, project, "lookup")
request = _verify_protobuf_call(
http,
uri,
datastore_pb2.LookupRequest(),
retry=retry,
timeout=timeout,
)
self.assertEqual(list(request.keys), [key_pb1._pb, key_pb2._pb])
self.assertEqual(request.read_options, read_options._pb)
def _run_query_helper(
self,
read_consistency=None,
transaction=None,
namespace=None,
found=0,
retry=None,
timeout=None,
):
from google.cloud.datastore_v1.types import datastore as datastore_pb2
from google.cloud.datastore_v1.types import entity as entity_pb2
from google.cloud.datastore_v1.types import query as query_pb2
project = "PROJECT"
kind = "Nonesuch"
query_pb = self._make_query_pb(kind)
partition_kw = {"project_id": project}
if namespace is not None:
partition_kw["namespace_id"] = namespace
partition_id = entity_pb2.PartitionId(**partition_kw)
options_kw = {}
if read_consistency is not None:
options_kw["read_consistency"] = read_consistency
if transaction is not None:
options_kw["transaction"] = transaction
read_options = datastore_pb2.ReadOptions(**options_kw)
cursor = b"\x00"
batch_kw = {
"entity_result_type":
query_pb2.EntityResult.ResultType.FULL,
"end_cursor":
cursor,
"more_results":
query_pb2.QueryResultBatch.MoreResultsType.NO_MORE_RESULTS,
}
if found:
batch_kw["entity_results"] = [
query_pb2.EntityResult(entity=entity_pb2.Entity())
] * found
rsp_pb = datastore_pb2.RunQueryResponse(
batch=query_pb2.QueryResultBatch(**batch_kw))
http = _make_requests_session(
[_make_response(content=rsp_pb._pb.SerializeToString())])
client_info = _make_client_info()
client = mock.Mock(
_http=http,
_base_url="test.invalid",
_client_info=client_info,
spec=["_http", "_base_url", "_client_info"],
)
ds_api = self._make_one(client)
request = {
"project_id": project,
"partition_id": partition_id,
"read_options": read_options,
"query": query_pb,
}
kwargs = _make_retry_timeout_kwargs(retry, timeout, http)
response = ds_api.run_query(request=request, **kwargs)
self.assertEqual(response, rsp_pb._pb)
uri = _build_expected_url(client._base_url, project, "runQuery")
request = _verify_protobuf_call(
http,
uri,
datastore_pb2.RunQueryRequest(),
retry=retry,
timeout=timeout,
)
self.assertEqual(request.partition_id, partition_id._pb)
self.assertEqual(request.query, query_pb._pb)
self.assertEqual(request.read_options, read_options._pb)
