def test_allocate_ids_empty(self):
from google.cloud.datastore_v1.proto import datastore_pb2
project = "PROJECT"
rsp_pb = datastore_pb2.AllocateIdsResponse()
# Create mock HTTP and client with response.
http = _make_requests_session(
[_make_response(content=rsp_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.allocate_ids(project, [])
# Check the result and verify the callers.
self.assertEqual(response, rsp_pb)
self.assertEqual(list(response.keys), [])
uri = _build_expected_url(client._base_url, project, "allocateIds")
request = _verify_protobuf_call(http, uri, datastore_pb2.AllocateIdsRequest())
self.assertEqual(list(request.keys), [])
def test_allocate_ids(self):
# Setup Expected Response
expected_response = {}
expected_response = datastore_pb2.AllocateIdsResponse(
**expected_response)
# Mock the API response
channel = ChannelStub(responses=[expected_response])
client = datastore_v1.DatastoreClient(channel=channel)
# Setup Request
project_id = 'projectId-1969970175'
keys = []
response = client.allocate_ids(project_id, keys)
assert expected_response == response
assert len(channel.requests) == 1
expected_request = datastore_pb2.AllocateIdsRequest(
project_id=project_id, keys=keys)
actual_request = channel.requests[0][1]
assert expected_request == actual_request
def test_allocate_ids_non_empty(self):
from google.cloud.datastore_v1.proto import datastore_pb2
project = "PROJECT"
before_key_pbs = [
_make_key_pb(project, id_=None),
_make_key_pb(project, id_=None),
]
after_key_pbs = [_make_key_pb(project), _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])
# Create mock HTTP and client with response.
http = _make_requests_session(
[_make_response(content=rsp_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.allocate_ids(project, before_key_pbs)
# Check the result and verify the callers.
self.assertEqual(list(response.keys), after_key_pbs)
self.assertEqual(response, rsp_pb)
uri = _build_expected_url(client._base_url, project, "allocateIds")
request = _verify_protobuf_call(http, uri, datastore_pb2.AllocateIdsRequest())
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)
