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)
Esempio n. 6
0
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())
Esempio n. 8
0
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), {})
Esempio n. 10
0
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)
Esempio n. 11
0
    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)
Esempio n. 12
0
    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)
Esempio n. 13
0
    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)})
Esempio n. 14
0
    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)
Esempio n. 15
0
    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)]))
Esempio n. 16
0
    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)
Esempio n. 17
0
    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)
Esempio n. 18
0
    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"], [])
Esempio n. 19
0
    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)
Esempio n. 20
0
    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)
Esempio n. 21
0
    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)
Esempio n. 22
0
    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)
Esempio n. 23
0
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
Esempio n. 24
0
    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)
Esempio n. 25
0
    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)
Esempio n. 26
0
    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)