def test_no_datastore_or_global_cache():
def MockEntity(*path):
key = ds_key_module.Key(*path, project="testing")
return entity.Entity(key=key)
mock_entity = MockEntity("what", "ever")
with pytest.raises(TypeError):
_api.put(mock_entity,
_options.Options(use_datastore=False)).result()
def test__use_datastore_from_options(self):
class SomeKind(model.Model):
pass
context = self._make_one()
with context.use():
key = key_module.Key("SomeKind", 1)
options = _options.Options(use_datastore=False)
assert context._use_datastore(key, options) is False
def test_delete_in_transaction(_datastore_api, in_context):
future = tasklets.Future()
_datastore_api.delete.return_value = future
with in_context.new(transaction=b"tx123").use():
key = key_module.Key("a", "b", app="c")
assert key.delete() is None
_datastore_api.delete.assert_called_once_with(
key._key, _options.Options())
def test_allocate_ids_callback(_datastore_allocate_ids):
options = _options.Options()
batch = _api._AllocateIdsBatch(options)
batch.futures = futures = [tasklets.Future(), tasklets.Future()]
rpc = utils.future_result(
mock.Mock(keys=["key1", "key2"], spec=("key", )))
batch.allocate_ids_callback(rpc)
results = [future.result() for future in futures]
assert results == ["key1", "key2"]
def test_allocate_ids_callback_w_exception(_datastore_allocate_ids):
options = _options.Options()
batch = _api._AllocateIdsBatch(options)
batch.futures = futures = [tasklets.Future(), tasklets.Future()]
error = Exception("spurious error")
rpc = tasklets.Future()
rpc.set_exception(error)
batch.allocate_ids_callback(rpc)
assert [future.exception() for future in futures] == [error, error]
def test__global_cache_timeout_from_options(self):
class SomeKind(model.Model):
pass
context = self._make_one()
with context.use():
key = "whocares"
options = _options.Options(global_cache_timeout=49)
assert context._global_cache_timeout(key, options) == 49
def test__use_global_cache_from_options(self):
class SomeKind(model.Model):
pass
context = self._make_one(global_cache="yes, there is one")
with context.use():
key = "whocares"
options = _options.Options(use_global_cache=False)
assert context._use_global_cache(key, options=options) is False
def test__use_cache_default_policy(self):
class SomeKind(model.Model):
pass
context = self._make_one()
with context.use():
key = key_module.Key("SomeKind", 1)
options = _options.Options()
assert context._use_cache(key, options) is True
def prepare_to_commit(transaction):
"""Signal that we're ready to commit a transaction.
Currently just used to signal to the commit batch that we're not going to
need to call `AllocateIds`, because we're ready to commit now.
Args:
transaction (bytes): The transaction id about to be committed.
"""
batch = _get_commit_batch(transaction, _options.Options())
batch.preparing_to_commit = True
def test_cache_enabled(Batch, global_cache):
key = key_module.Key("SomeKind", 1)
cache_key = _cache.global_cache_key(key._key)
batch = Batch.return_value
batch.delete.return_value = future_result(None)
future = _api.delete(key._key, _options.Options())
assert future.result() is None
assert global_cache.get([cache_key]) == [None]
def test_without_datastore(Batch, global_cache):
key = key_module.Key("SomeKind", 1)
cache_key = _cache.global_cache_key(key._key)
global_cache.set({cache_key: b"foo"})
batch = Batch.return_value
batch.delete.side_effect = Exception("Shouldn't use Datastore")
future = _api.delete(key._key, _options.Options(use_datastore=False))
assert future.result() is None
assert global_cache.get([cache_key]) == [None]
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_no_datastore_incomplete_key(Batch, global_cache):
class SomeKind(model.Model):
pass
key = key_module.Key("SomeKind", None)
entity = SomeKind(key=key)
future = _api.put(
model._entity_to_ds_entity(entity),
_options.Options(use_datastore=False),
)
with pytest.raises(TypeError):
future.result()
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 test_no_transaction(datastore_pb2, in_context):
class Mutation:
def __init__(self, delete=None):
self.delete = delete
def __eq__(self, other):
return self.delete == other.delete
datastore_pb2.Mutation = Mutation
key1 = key_module.Key("SomeKind", 1)._key
key2 = key_module.Key("SomeKind", 2)._key
key3 = key_module.Key("SomeKind", 3)._key
_api.delete(key1, _options.Options())
_api.delete(key2, _options.Options())
_api.delete(key3, _options.Options())
batch = in_context.batches[_api._NonTransactionalCommitBatch][()]
assert batch.mutations == [
Mutation(delete=key1.to_protobuf()),
Mutation(delete=key2.to_protobuf()),
Mutation(delete=key3.to_protobuf()),
]
def test_w_transaction(datastore_pb2, in_context):
class Mutation:
def __init__(self, delete=None):
self.delete = delete
def __eq__(self, other):
return self.delete == other.delete
with in_context.new(transaction=b"tx123").use() as context:
datastore_pb2.Mutation = Mutation
key1 = key_module.Key("SomeKind", 1)._key
key2 = key_module.Key("SomeKind", 2)._key
key3 = key_module.Key("SomeKind", 3)._key
assert _api.delete(key1, _options.Options()).result() is None
assert _api.delete(key2, _options.Options()).result() is None
assert _api.delete(key3, _options.Options()).result() is None
batch = context.commit_batches[b"tx123"]
assert batch.mutations == [
Mutation(delete=key1.to_protobuf()),
Mutation(delete=key2.to_protobuf()),
Mutation(delete=key3.to_protobuf()),
]
def test_no_key_returned(Batch, global_cache):
class SomeKind(model.Model):
pass
key = key_module.Key("SomeKind", 1)
cache_key = _cache.global_cache_key(key._key)
entity = SomeKind(key=key)
batch = Batch.return_value
batch.put.return_value = future_result(None)
future = _api.put(model._entity_to_ds_entity(entity), _options.Options())
assert future.result() is None
assert global_cache.get([cache_key]) == [None]
def test_idle_callback(_datastore_allocate_ids):
options = _options.Options()
batch = _api._AllocateIdsBatch(options)
batch.add([
key_module.Key("SomeKind", None)._key,
key_module.Key("SomeKind", None)._key,
])
key_pbs = [key.to_protobuf() for key in batch.keys]
batch.idle_callback()
_datastore_allocate_ids.assert_called_once_with(key_pbs,
retries=None,
timeout=None)
rpc = _datastore_allocate_ids.return_value
rpc.add_done_callback.assert_called_once_with(
batch.allocate_ids_callback)
def test_no_transaction(datastore_pb2, in_context):
class Mutation:
def __init__(self, upsert=None):
self.upsert = upsert
def __eq__(self, other):
return self.upsert is other.upsert
eventloop = mock.Mock(spec=("add_idle", "run"))
with in_context.new(eventloop=eventloop).use() as context:
datastore_pb2.Mutation = Mutation
entity1, entity2, entity3 = object(), object(), object()
future1 = _api.put(entity1, _options.Options())
future2 = _api.put(entity2, _options.Options())
future3 = _api.put(entity3, _options.Options())
batch = context.batches[_api._NonTransactionalCommitBatch][()]
assert batch.mutations == [
Mutation(upsert=entity1),
Mutation(upsert=entity2),
Mutation(upsert=entity3),
]
assert batch.futures == [future1, future2, future3]
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 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 commit(transaction, retries=None, timeout=None):
"""Commit a transaction.
Args:
transaction (bytes): The transaction id to commit.
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`.
Returns:
tasklets.Future: Result will be none, will finish when the transaction
is committed.
"""
batch = _get_commit_batch(transaction, _options.Options())
return batch.commit(retries=retries, timeout=timeout)
def test_idle_callback(_datastore_commit, _process_commit, context):
eventloop = mock.Mock(spec=("queue_rpc", "run"))
rpc = tasklets.Future("_datastore_commit")
_datastore_commit.return_value = rpc
with context.new(eventloop=eventloop).use() as context:
mutation1, mutation2 = object(), object()
batch = _api._NonTransactionalCommitBatch(_options.Options())
batch.mutations = [mutation1, mutation2]
batch.idle_callback()
_datastore_commit.assert_called_once_with(
[mutation1, mutation2], None, retries=None, timeout=None
)
rpc.set_result(None)
_process_commit.assert_called_once_with(rpc, batch.futures)
def test_no_datastore(Batch, global_cache):
class SomeKind(model.Model):
pass
key = key_module.Key("SomeKind", 1)
cache_key = _cache.global_cache_key(key._key)
entity = SomeKind(key=key)
cache_value = model._entity_to_protobuf(entity).SerializeToString()
batch = Batch.return_value
batch.put.return_value = future_result(None)
future = _api.put(
model._entity_to_ds_entity(entity),
_options.Options(use_datastore=False),
)
assert future.result() is None
assert global_cache.get([cache_key]) == [cache_value]
def test_commit(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
)
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_bad_option():
with pytest.raises(NotImplementedError):
_api._get_commit_batch(b"123", _options.Options(retries=5))
def test_commit(get_commit_batch):
_api.commit(b"123")
get_commit_batch.assert_called_once_with(b"123", _options.Options())
get_commit_batch.return_value.commit.assert_called_once_with(
retries=None, timeout=None
)
def test_constructor():
options = _options.Options()
batch = _api._AllocateIdsBatch(options)
assert batch.options is options
assert batch.keys == []
assert batch.futures == []
def test_add():
options = _options.Options()
batch = _api._AllocateIdsBatch(options)
futures = batch.add(["key1", "key2"])
assert batch.keys == ["key1", "key2"]
assert batch.futures == futures
