def test_allocate_ids_non_empty(self):
from google.cloud.grpc.datastore.v1 import datastore_pb2
PROJECT = 'PROJECT'
before_key_pbs = [
self._make_key_pb(PROJECT, id_=None),
self._make_key_pb(PROJECT, id_=None),
]
after_key_pbs = [
self._make_key_pb(PROJECT),
self._make_key_pb(PROJECT, id_=2345),
]
rsp_pb = datastore_pb2.AllocateIdsResponse()
rsp_pb.keys.add().CopyFrom(after_key_pbs[0])
rsp_pb.keys.add().CopyFrom(after_key_pbs[1])
conn = self._make_one()
URI = '/'.join([
conn.api_base_url,
conn.API_VERSION,
'projects',
PROJECT + ':allocateIds',
])
http = conn._http = Http({'status': '200'}, rsp_pb.SerializeToString())
self.assertEqual(conn.allocate_ids(PROJECT, before_key_pbs),
after_key_pbs)
cw = http._called_with
self._verifyProtobufCall(cw, URI, conn)
rq_class = datastore_pb2.AllocateIdsRequest
request = rq_class()
request.ParseFromString(cw['body'])
self.assertEqual(len(request.keys), len(before_key_pbs))
for key_before, key_after in zip(before_key_pbs, request.keys):
self.assertEqual(key_before, key_after)
def test_allocate_ids_empty(self):
from google.cloud.grpc.datastore.v1 import datastore_pb2
PROJECT = 'PROJECT'
rsp_pb = datastore_pb2.AllocateIdsResponse()
conn = self._make_one()
URI = '/'.join([
conn.api_base_url,
conn.API_VERSION,
'projects',
PROJECT + ':allocateIds',
])
http = conn._http = Http({'status': '200'}, rsp_pb.SerializeToString())
self.assertEqual(conn.allocate_ids(PROJECT, []), [])
cw = http._called_with
self._verifyProtobufCall(cw, URI, conn)
rq_class = datastore_pb2.AllocateIdsRequest
request = rq_class()
request.ParseFromString(cw['body'])
self.assertEqual(list(request.keys), [])