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