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