def test_idle_callback_failure(datastore_allocate_ids, in_context):
def Mutation():
path = [entity_pb2.Key.PathElement(kind="SomeKind")]
return datastore_pb2.Mutation(
upsert=entity_pb2.Entity(key=entity_pb2.Key(path=path))
)
mutation1, mutation2 = Mutation(), Mutation()
batch = _api._TransactionalCommitBatch(b"123", _options.Options())
batch.incomplete_mutations = [mutation1, mutation2]
future1, future2 = tasklets.Future(), tasklets.Future()
batch.incomplete_futures = [future1, future2]
rpc = tasklets.Future("_datastore_allocate_ids")
datastore_allocate_ids.return_value = rpc
eventloop = mock.Mock(spec=("queue_rpc", "run"))
with in_context.new(eventloop=eventloop).use():
batch.idle_callback()
error = Exception("Spurious error")
rpc.set_exception(error)
allocating_ids = batch.allocating_ids[0]
assert future1.exception() is error
assert future2.exception() is error
assert allocating_ids.result() is None
def test_commit_allocating_ids(
datastore_commit, process_commit, in_context
):
batch = _api._TransactionalCommitBatch(b"123", _options.Options())
batch.futures = object()
batch.mutations = object()
batch.transaction = b"abc"
allocated_ids = tasklets.Future("Already allocated ids")
allocated_ids.set_result(None)
batch.allocating_ids.append(allocated_ids)
allocating_ids = tasklets.Future("AllocateIds")
batch.allocating_ids.append(allocating_ids)
rpc = tasklets.Future("_datastore_commit")
datastore_commit.return_value = rpc
eventloop = mock.Mock(spec=("queue_rpc", "run", "call_soon"))
eventloop.call_soon = lambda f, *args, **kwargs: f(*args, **kwargs)
with in_context.new(eventloop=eventloop).use():
future = batch.commit()
datastore_commit.assert_not_called()
process_commit.assert_not_called()
allocating_ids.set_result(None)
datastore_commit.assert_called_once_with(
batch.mutations, transaction=b"abc", retries=None, timeout=None
)
rpc.set_result(None)
process_commit.assert_called_once_with(rpc, batch.futures)
assert future.result() is None
def test_commit_error(datastore_commit, process_commit, in_context):
batch = _api._TransactionalCommitBatch(b"123", _options.Options())
batch.futures = object()
batch.mutations = object()
batch.transaction = b"abc"
rpc = tasklets.Future("_datastore_commit")
datastore_commit.return_value = rpc
eventloop = mock.Mock(spec=("queue_rpc", "run", "call_soon"))
eventloop.call_soon = lambda f, *args, **kwargs: f(*args, **kwargs)
with in_context.new(eventloop=eventloop).use():
future = batch.commit()
datastore_commit.assert_called_once_with(batch.mutations,
transaction=b"abc",
retries=None,
timeout=None)
error = Exception("Spurious error")
rpc.set_exception(error)
process_commit.assert_called_once_with(rpc, batch.futures)
assert future.exception() is error
def test_idle_callback_success(datastore_allocate_ids, in_context):
def Mutation():
path = [entity_pb2.Key.PathElement(kind="SomeKind")]
return datastore_pb2.Mutation(upsert=entity_pb2.Entity(
key=entity_pb2.Key(path=path)))
mutation1, mutation2 = Mutation(), Mutation()
batch = _api._TransactionalCommitBatch(b"123", _options.Options())
batch.incomplete_mutations = [mutation1, mutation2]
future1, future2 = tasklets.Future(), tasklets.Future()
batch.incomplete_futures = [future1, future2]
rpc = tasklets.Future("_datastore_allocate_ids")
datastore_allocate_ids.return_value = rpc
eventloop = mock.Mock(spec=("queue_rpc", "run"))
with in_context.new(eventloop=eventloop).use():
batch.idle_callback()
rpc.set_result(
mock.Mock(keys=[
entity_pb2.Key(path=[
entity_pb2.Key.PathElement(kind="SomeKind", id=1)
]),
entity_pb2.Key(path=[
entity_pb2.Key.PathElement(kind="SomeKind", id=2)
]),
]))
allocating_ids = batch.allocating_ids[0]
assert future1.result().path[0].id == 1
assert mutation1.upsert.key.path[0].id == 1
assert future2.result().path[0].id == 2
assert mutation2.upsert.key.path[0].id == 2
assert allocating_ids.result() is None
def test_commit_nothing_to_do(in_context):
batch = _api._TransactionalCommitBatch({})
eventloop = mock.Mock(spec=("queue_rpc", "run"))
with in_context.new(eventloop=eventloop).use():
future = batch.commit()
eventloop.queue_rpc.assert_not_called()
assert future.result() is None
def test_commit(datastore_commit, process_commit, in_context):
batch = _api._TransactionalCommitBatch({})
batch.futures = object()
batch.mutations = object()
batch.transaction = b"abc"
rpc = tasklets.Future("_datastore_commit")
datastore_commit.return_value = rpc
eventloop = mock.Mock(spec=("queue_rpc", "run"))
with in_context.new(eventloop=eventloop).use():
future = batch.commit()
datastore_commit.assert_called_once_with(batch.mutations,
transaction=b"abc",
retries=None)
rpc.set_result(None)
process_commit.assert_called_once_with(rpc, batch.futures)
assert future.result() is None
def test_idle_callback_nothing_to_do():
batch = _api._TransactionalCommitBatch(b"123", _options.Options())
batch.idle_callback()
assert not batch.allocating_ids
def test_idle_callback_nothing_to_do():
batch = _api._TransactionalCommitBatch({})
batch.idle_callback()
assert not batch.allocating_ids