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()