def test_commit_callback_exception():
future1, future2 = tasklets.Future(), tasklets.Future()
batch = _api._CommitBatch({})
batch.futures = [future1, future2]
error = Exception("Spurious error.")
rpc = tasklets.Future()
rpc.set_exception(error)
batch.commit_callback(rpc)
assert future1.exception() is error
assert future2.exception() is error
def test_idle_callback_no_transaction(_datastore_commit, runstate):
runstate.eventloop = mock.Mock(spec=("queue_rpc", "run"))
mutation1, mutation2 = object(), object()
batch = _api._CommitBatch({})
batch.mutations = [mutation1, mutation2]
batch.idle_callback()
rpc = _datastore_commit.return_value
_datastore_commit.assert_called_once_with([mutation1, mutation2], None)
runstate.eventloop.queue_rpc.assert_called_once_with(
rpc, batch.commit_callback
)
def test_idle_callback_w_transaction(_datastore_commit, context):
eventloop = mock.Mock(spec=("queue_rpc", "run"))
with context.new(eventloop=eventloop) as context:
mutation1, mutation2 = object(), object()
batch = _api._CommitBatch({"transaction": b"tx123"})
batch.mutations = [mutation1, mutation2]
batch.idle_callback()
rpc = _datastore_commit.return_value
_datastore_commit.assert_called_once_with(
[mutation1, mutation2], b"tx123"
)
context.eventloop.queue_rpc.assert_called_once_with(
rpc, batch.commit_callback
)
def test_commit_callback():
future1, future2 = tasklets.Future(), tasklets.Future()
batch = _api._CommitBatch({})
batch.futures = [future1, future2]
key1 = mock.Mock(path=["one", "two"], spec=("path",))
mutation1 = mock.Mock(key=key1, spec=("key",))
key2 = mock.Mock(path=[], spec=("path",))
mutation2 = mock.Mock(key=key2, spec=("key",))
response = mock.Mock(
mutation_results=(mutation1, mutation2), spec=("mutation_results",)
)
rpc = tasklets.Future()
rpc.set_result(response)
batch.commit_callback(rpc)
assert future1.result() is key1
assert future2.result() is None