def test_delete_async(_datastore_api):
key = key_module.Key("a", "b", app="c")
future = tasklets.Future()
_datastore_api.delete.return_value = future
future.set_result("result")
result = key.delete_async().get_result()
_datastore_api.delete.assert_called_once_with(
key._key, _options.Options()
)
assert result == "result"
def add(self, key):
"""Add a key to delete from the cache.
Arguments:
key (bytes): The key to delete.
Returns:
tasklets.Future: Eventual result will be ``None``.
"""
future = tasklets.Future(info=self.future_info(key))
self.keys.append(key)
self.futures.append(future)
return future
def test__datastore_allocate_ids(stub, datastore_pb2):
keys = object()
api = stub.return_value
future = tasklets.Future()
future.set_result("response")
api.AllocateIds.future.return_value = future
assert _api._datastore_allocate_ids(keys).result() == "response"
datastore_pb2.AllocateIdsRequest.assert_called_once_with(
project_id="testing", keys=keys)
request = datastore_pb2.AllocateIdsRequest.return_value
assert api.AllocateIds.future.called_once_with(request)
def test_delete(_datastore_api):
class Simple(model.Model):
pass
future = tasklets.Future()
_datastore_api.delete.return_value = future
future.set_result("result")
key = key_module.Key("Simple", "b", app="c")
assert key.delete() == "result"
_datastore_api.delete.assert_called_once_with(
key._key, _options.Options()
)
def test__advance_tasklet_dependency_raises():
def generator_function(dependency):
yield dependency
error = Exception("Spurious error.")
dependency = tasklets.Future()
generator = generator_function(dependency)
future = tasklets._TaskletFuture(generator)
future._advance_tasklet()
dependency.set_exception(error)
assert future.exception() is error
with pytest.raises(Exception):
future.result()
def add(self, key):
"""Add a key to the batch to look up.
Args:
key (datastore.Key): The key to look up.
Returns:
tasklets.Future: A future for the eventual result.
"""
todo_key = key.to_protobuf().SerializeToString()
future = tasklets.Future(info="Lookup({})".format(key))
self.todo.setdefault(todo_key, []).append(future)
return future
def test_error(_datastore_api):
error = Exception("Spurious error.")
def callback():
raise error
begin_future = tasklets.Future("begin transaction")
_datastore_api.begin_transaction.return_value = begin_future
rollback_future = tasklets.Future("rollback transaction")
_datastore_api.rollback.return_value = rollback_future
future = _transaction.transaction_async(callback)
_datastore_api.begin_transaction.assert_called_once_with(False,
retries=0)
begin_future.set_result(b"tx123")
_datastore_api.rollback.assert_called_once_with(b"tx123")
rollback_future.set_result(None)
assert future.exception() is error
def test_propagation_independent_already_in_transaction(
_datastore_api, in_context):
def callback():
return "I tried, momma."
begin_future = tasklets.Future("begin transaction")
_datastore_api.begin_transaction.return_value = begin_future
commit_future = tasklets.Future("commit transaction")
_datastore_api.commit.return_value = commit_future
with mock.patch(
"google.cloud.ndb._transaction.transaction_async_",
side_effect=_transaction.transaction_async_,
) as transaction_async_:
with in_context.new(transaction=b"tx123").use():
future = _transaction.transaction_async(
callback,
join=True,
propagation=context_module.TransactionOptions.INDEPENDENT,
)
_datastore_api.begin_transaction.assert_called_once_with(False,
retries=0)
begin_future.set_result(b"tx456")
_datastore_api.commit.assert_called_once_with(b"tx456", retries=0)
commit_future.set_result(None)
assert future.result() == "I tried, momma."
transaction_async_.assert_called_once_with(
callback,
3,
False,
False,
True,
None,
)
def test_cancel_waiting_on_dependency(in_context):
def generator_function(dependency):
yield dependency
dependency = tasklets.Future()
generator = generator_function(dependency)
future = tasklets._TaskletFuture(generator, in_context)
future._advance_tasklet()
future.cancel()
assert dependency.cancelled()
with pytest.raises(exceptions.Cancelled):
future.result()
def test_transient_error(core_retry, sleep):
core_retry.exponential_sleep_generator.return_value = itertools.count()
core_retry.if_transient_error.return_value = True
sleep_future = tasklets.Future("sleep")
sleep.return_value = sleep_future
callback = mock.Mock(side_effect=[Exception("Spurious error."), "foo"])
retry = _retry.retry_async(callback)
sleep_future.set_result(None)
assert retry().result() == "foo"
sleep.assert_called_once_with(0)
def test_get_async(_entity_from_protobuf, _datastore_api):
ds_future = tasklets.Future()
_datastore_api.lookup.return_value = ds_future
_entity_from_protobuf.return_value = "the entity"
key = key_module.Key("a", "b", app="c")
future = key.get_async()
ds_future.set_result("ds_entity")
assert future.result() == "the entity"
_datastore_api.lookup.assert_called_once_with(key._key,
_options.ReadOptions())
_entity_from_protobuf.assert_called_once_with("ds_entity")
def test_explicit_retries(stub, _retry):
api = stub.return_value
future = tasklets.Future()
api.foo.future.return_value = future
_retry.retry_async.return_value = mock.Mock(return_value=future)
future.set_result("bar")
request = object()
assert _api.make_call("foo", request, retries=4).result() == "bar"
_retry.retry_async.assert_called_once()
tasklet = _retry.retry_async.call_args[0][0]
assert tasklet().result() == "bar"
retries = _retry.retry_async.call_args[1]["retries"]
assert retries == 4
def test_success_callback_is_tasklet(_datastore_api):
tasklet = tasklets.Future("tasklet")
def callback():
return tasklet
begin_future = tasklets.Future("begin transaction")
_datastore_api.begin_transaction.return_value = begin_future
commit_future = tasklets.Future("commit transaction")
_datastore_api.commit.return_value = commit_future
future = _transaction.transaction_async(callback)
_datastore_api.begin_transaction.assert_called_once_with(False, retries=0)
begin_future.set_result(b"tx123")
tasklet.set_result("I tried, momma.")
_datastore_api.commit.assert_called_once_with(b"tx123", retries=0)
commit_future.set_result(None)
assert future.result() == "I tried, momma."
def add(self, key, value):
"""Add a key, value pair to store in the cache.
Arguments:
key (bytes): The key to store in the cache.
value (bytes): The value to store in the cache.
Returns:
tasklets.Future: Eventual result will be ``None``.
"""
future = tasklets.Future(info=self.future_info(key, value))
self.todo[key] = value
self.futures.append(future)
return future
def test_propagation_allowed_not_yet_in_transaction(_datastore_api):
def callback():
return "I tried, momma."
begin_future = tasklets.Future("begin transaction")
_datastore_api.begin_transaction.return_value = begin_future
commit_future = tasklets.Future("commit transaction")
_datastore_api.commit.return_value = commit_future
with mock.patch(
"google.cloud.ndb._transaction.transaction_async_",
side_effect=_transaction.transaction_async_,
) as transaction_async_:
future = _transaction.transaction_async(
callback,
join=False,
propagation=context_module.TransactionOptions.ALLOWED,
)
_datastore_api.begin_transaction.assert_called_once_with(False,
retries=0)
begin_future.set_result(b"tx123")
_datastore_api.commit.assert_called_once_with(b"tx123", retries=0)
commit_future.set_result(None)
assert future.result() == "I tried, momma."
transaction_async_.assert_called_once_with(
callback,
3,
False,
True,
True,
None,
)
def test_check_success_failure(_eventloop):
error = Exception("Spurious error")
def side_effects(future):
yield
yield
future.set_exception(error)
yield
future = tasklets.Future()
_eventloop.run1.side_effect = side_effects(future)
with pytest.raises(Exception) as error_context:
future.check_success()
assert error_context.value is error
def delete(self, key):
"""Add a key to batch to be deleted.
Args:
entity_pb (datastore.Key): The entity's key to be deleted.
Returns:
tasklets.Future: Result will be :data:`None`, always.
"""
key_pb = key.to_protobuf()
future = tasklets.Future(info="delete({})".format(key_pb))
mutation = datastore_pb2.Mutation(delete=key_pb)
self.mutations.append(mutation)
self.futures.append(future)
return future
def put(self, entity_pb):
"""Add an entity to batch to be stored.
Args:
entity_pb (datastore_v1.types.Entity): The entity to be stored.
Returns:
tasklets.Future: Result will be completed datastore key
(entity_pb2.Key) for the entity.
"""
future = tasklets.Future(info="put({})".format(entity_pb))
mutation = datastore_pb2.Mutation(upsert=entity_pb)
self.mutations.append(mutation)
self.futures.append(future)
return future
def test_found_missing_deferred(context):
def key_pb(key):
mock_key = mock.Mock(spec=("SerializeToString", ))
mock_key.SerializeToString.return_value = key
return mock_key
eventloop = mock.Mock(spec=("add_idle", "run"))
with context.new(eventloop=eventloop).use() as context:
future1, future2, future3 = (tasklets.Future() for _ in range(3))
batch = _api._LookupBatch(_options.ReadOptions())
batch.todo.update({
"foo": [future1],
"bar": [future2],
"baz": [future3]
})
entity1 = mock.Mock(key=key_pb("foo"), spec=("key", ))
entity2 = mock.Mock(key=key_pb("bar"), spec=("key", ))
response = mock.Mock(
found=[mock.Mock(entity=entity1, spec=("entity", ))],
missing=[mock.Mock(entity=entity2, spec=("entity", ))],
deferred=[key_pb("baz")],
spec=("found", "missing", "deferred"),
)
rpc = tasklets.Future()
rpc.set_result(response)
batch.lookup_callback(rpc)
assert future1.result() is entity1
assert future2.result() is _api._NOT_FOUND
assert future3.running()
next_batch = context.batches[_api._LookupBatch][()]
assert next_batch.todo == {"baz": [future3]}
assert context.eventloop.add_idle.call_count == 1
def test_found_missing_deferred(runstate):
runstate.eventloop = mock.Mock(spec=("add_idle", "run"))
future1, future2, future3 = (tasklets.Future() for _ in range(3))
batch = {"foo": [future1], "bar": [future2], "baz": [future3]}
entity1 = mock.Mock(key="foo", spec=("key", ))
entity2 = mock.Mock(key="bar", spec=("key", ))
response = mock.Mock(
found=[mock.Mock(entity=entity1, spec=("entity", ))],
missing=[mock.Mock(entity=entity2, spec=("entity", ))],
deferred=["baz"],
spec=("found", "missing", "deferred"),
)
rpc = tasklets.Future()
rpc.set_result(response)
callback = _api.BatchLookupCallback(batch)
callback(rpc)
assert future1.result() is entity1
assert future2.result() is _api._NOT_FOUND
assert future3.running()
assert runstate.batches[_api._BATCH_LOOKUP] == {"baz": [future3]}
runstate.eventloop.add_idle.assert_called_once_with(
_api._perform_batch_lookup)
def test_success(_datastore_api):
on_commit_callback = mock.Mock()
def callback():
context_module.get_context().call_on_commit(on_commit_callback)
return "I tried, momma."
begin_future = tasklets.Future("begin transaction")
_datastore_api.begin_transaction.return_value = begin_future
commit_future = tasklets.Future("commit transaction")
_datastore_api.commit.return_value = commit_future
future = _transaction.transaction_async(callback)
_datastore_api.begin_transaction.assert_called_once_with(False,
retries=0)
begin_future.set_result(b"tx123")
_datastore_api.commit.assert_called_once_with(b"tx123", retries=0)
commit_future.set_result(None)
assert future.result() == "I tried, momma."
on_commit_callback.assert_called_once_with()
def test__datastore_lookup(datastore_pb2, context):
client = mock.Mock(project="theproject", spec=("project", ))
stub = mock.Mock(spec=("Lookup", ))
with context.new(client=client, stub=stub).use() as context:
context.stub.Lookup = Lookup = mock.Mock(spec=("future", ))
future = tasklets.Future()
future.set_result("response")
Lookup.future.return_value = future
assert (_api._datastore_lookup(["foo", "bar"],
None).result() == "response")
datastore_pb2.LookupRequest.assert_called_once_with(
project_id="theproject", keys=["foo", "bar"], read_options=None)
context.stub.Lookup.future.assert_called_once_with(
datastore_pb2.LookupRequest.return_value)
def test_get_with_cache_miss(_entity_from_protobuf, _datastore_api):
class Simple(model.Model):
pass
ds_future = tasklets.Future()
ds_future.set_result("ds_entity")
_datastore_api.lookup.return_value = ds_future
_entity_from_protobuf.return_value = "the entity"
key = key_module.Key("Simple", "b", app="c")
assert key.get(use_cache=True) == "the entity"
_datastore_api.lookup.assert_called_once_with(
key._key, _options.ReadOptions(use_cache=True))
_entity_from_protobuf.assert_called_once_with("ds_entity")
def commit(self, retries=None, timeout=None):
"""Commit transaction.
Args:
retries (int): Number of times to potentially retry the call. If
:data:`None` is passed, will use
:data:`_retry._DEFAULT_RETRIES`. If :data:`0` is passed, the
call is attempted only once.
timeout (float): Timeout, in seconds, to pass to gRPC call. If
:data:`None` is passed, will use :data:`_DEFAULT_TIMEOUT`.
"""
# It's tempting to do something like:
#
# if not self.mutations:
# return
#
# However, even if there are no mutations to save, we still need to
# send a COMMIT to the Datastore. It would appear that failing to do so
# will make subsequent writes hang indefinitely as Datastore apparently
# achieves consistency during a transaction by preventing writes.
# Wait for any calls to AllocateIds that have been fired off so we
# don't allocate ids again in the commit.
for future in self.allocating_ids:
if not future.done():
yield future
future = tasklets.Future("Commit")
futures = self.futures
def commit_callback(rpc):
_process_commit(rpc, futures)
exception = rpc.exception()
if exception:
future.set_exception(exception)
else:
future.set_result(None)
rpc = _datastore_commit(
self.mutations,
transaction=self.transaction,
retries=retries,
timeout=timeout,
)
rpc.add_done_callback(commit_callback)
yield future
def test_read_write(stub, datastore_pb2):
api = stub.return_value
future = tasklets.Future()
future.set_result("response")
api.BeginTransaction.future.return_value = future
assert _api._datastore_begin_transaction(False).result() == "response"
datastore_pb2.TransactionOptions.assert_called_once_with(
read_write=datastore_pb2.TransactionOptions.ReadWrite())
transaction_options = datastore_pb2.TransactionOptions.return_value
datastore_pb2.BeginTransactionRequest.assert_called_once_with(
project_id="testing", transaction_options=transaction_options)
request = datastore_pb2.BeginTransactionRequest.return_value
assert api.BeginTransaction.future.called_once_with(request)
def test_other_error(stub):
api = stub.return_value
future = tasklets.Future()
api.foo.future.return_value = future
class DummyException(Exception):
pass
try:
raise DummyException("Have to raise in order to get traceback")
except Exception as error:
future.set_exception(error)
request = object()
with pytest.raises(DummyException):
_api.make_call("foo", request, retries=0).result()
def test_delete_no_cache(_datastore_api, in_context):
class Simple(model.Model):
pass
future = tasklets.Future()
_datastore_api.delete.return_value = future
future.set_result("result")
key = key_module.Key("Simple", "b", app="c")
mock_cached_entity = mock.Mock(_key=key)
in_context.cache[key] = mock_cached_entity
assert key.delete(use_cache=False) == "result"
assert in_context.cache[key] == mock_cached_entity
_datastore_api.delete.assert_called_once_with(
key._key, _options.Options(use_cache=False))
def add(self, key):
"""Add a key to get from the cache.
Arguments:
key (bytes): The key to get from the cache.
Returns:
tasklets.Future: Eventual result will be the entity retrieved from
the cache (``bytes``) or ``None``.
"""
future = tasklets.Future(info=self.future_info(key))
futures = self.todo.get(key)
if futures is None:
self.todo[key] = futures = []
self.keys.append(key)
futures.append(future)
return future
def test_get_no_cache(_entity_from_protobuf, _datastore_api, in_context):
class Simple(model.Model):
pass
ds_future = tasklets.Future()
ds_future.set_result("ds_entity")
_datastore_api.lookup.return_value = ds_future
_entity_from_protobuf.return_value = "the entity"
key = key_module.Key("Simple", "b", app="c")
mock_cached_entity = mock.Mock(_key=key)
in_context.cache[key] = mock_cached_entity
assert key.get(use_cache=False) == "the entity"
_datastore_api.lookup.assert_called_once_with(
key._key, _options.ReadOptions(use_cache=False))
_entity_from_protobuf.assert_called_once_with("ds_entity")
def test_get_with_cache_hit(_entity_from_protobuf, _datastore_api,
in_context):
class Simple(model.Model):
pass
ds_future = tasklets.Future()
ds_future.set_result("ds_entity")
_datastore_api.lookup.return_value = ds_future
_entity_from_protobuf.return_value = "the entity"
key = key_module.Key("Simple", "b", app="c")
mock_cached_entity = mock.Mock(_key=key)
in_context.cache[key] = mock_cached_entity
assert key.get(use_cache=True) == mock_cached_entity
_datastore_api.lookup.assert_not_called()
_entity_from_protobuf.assert_not_called()