def test_add_idle(EventLoop):
EventLoop.return_value = loop = unittest.mock.Mock(
spec=("run", "add_idle")
)
with _runstate.ndb_context():
_eventloop.add_idle("foo", "bar", baz="qux")
loop.add_idle.assert_called_once_with("foo", "bar", baz="qux")
def test_it():
futures = [tasklets.Future() for _ in range(3)]
def callback():
futures[1].set_result(42)
_eventloop.add_idle(callback)
future = tasklets.wait_any(futures)
assert future is futures[1]
assert future.result() == 42
def test_toplevel():
@tasklets.toplevel
def generator_function(value):
future = tasklets.Future(value)
future.set_result(value)
x = yield future
raise tasklets.Return(x + 3)
idle = mock.Mock(__name__="idle", return_value=None)
_eventloop.add_idle(idle)
result = generator_function(8)
assert result == 11
idle.assert_called_once_with()
def get_batch(batch_cls, options=None):
"""Gets a data structure for storing batched calls to Datastore Lookup.
The batch data structure is stored in the current context. If there is
not already a batch started, a new structure is created and an idle
callback is added to the current event loop which will eventually perform
the batch look up.
Args:
batch_cls (type): Class representing the kind of operation being
batched.
options (_options.ReadOptions): The options for the request. Calls with
different options will be placed in different batches.
Returns:
batch_cls: An instance of the batch class.
"""
# prevent circular import in Python 2.7
from google.cloud.ndb import context as context_module
context = context_module.get_context()
batches = context.batches.get(batch_cls)
if batches is None:
context.batches[batch_cls] = batches = {}
if options is not None:
options_key = tuple(
sorted(((key, value) for key, value in options.items()
if value is not None)))
else:
options_key = ()
batch = batches.get(options_key)
if batch is not None and not batch.full():
return batch
def idler(batch):
def idle():
if batches.get(options_key) is batch:
del batches[options_key]
batch.idle_callback()
return idle
batches[options_key] = batch = batch_cls(options)
_eventloop.add_idle(idler(batch))
return batch
def _get_lookup_batch():
"""Gets a data structure for storing batched calls to Datastore Lookup.
The batch data structure is stored in the current run state. If there is
not already a batch started, a new structure is created and an idle
callback is added to the current event loop which will eventually perform
the batch look up.
Returns:
Dict[~datastore_v1.proto.entity_pb2.Key, List[~tasklets.Future]]
"""
state = _runstate.current()
batch = state.batches.get(_BATCH_LOOKUP)
if batch is not None:
return batch
state.batches[_BATCH_LOOKUP] = batch = {}
_eventloop.add_idle(_perform_batch_lookup)
return batch
def _get_batch(batch_cls, options):
"""Gets a data structure for storing batched calls to Datastore Lookup.
The batch data structure is stored in the current context. If there is
not already a batch started, a new structure is created and an idle
callback is added to the current event loop which will eventually perform
the batch look up.
Args:
batch_cls (type): Class representing the kind of operation being
batched.
options (_options.ReadOptions): The options for the request. Calls with
different options will be placed in different batches.
Returns:
batch_cls: An instance of the batch class.
"""
context = context_module.get_context()
batches = context.batches.get(batch_cls)
if batches is None:
context.batches[batch_cls] = batches = {}
options_key = tuple(
sorted(((key, value) for key, value in options.items()
if value is not None)))
batch = batches.get(options_key)
if batch is not None:
return batch
def idle():
batch = batches.pop(options_key)
batch.idle_callback()
batches[options_key] = batch = batch_cls(options)
_eventloop.add_idle(idle)
return batch
def test_it():
futures = [tasklets.Future() for _ in range(3)]
def make_callback(index, result):
def callback():
futures[index].set_result(result)
return callback
_eventloop.add_idle(make_callback(0, 42))
_eventloop.add_idle(make_callback(1, 43))
_eventloop.add_idle(make_callback(2, 44))
tasklets.wait_all(futures)
assert futures[0].done()
assert futures[0].result() == 42
assert futures[1].done()
assert futures[1].result() == 43
assert futures[2].done()
assert futures[2].result() == 44
def test_add_idle(context):
loop = mock.Mock(spec=("run", "add_idle"))
with context.new(eventloop=loop).use():
_eventloop.add_idle("foo", "bar", baz="qux")
loop.add_idle.assert_called_once_with("foo", "bar", baz="qux")
def test_add_idle(EventLoop):
EventLoop.return_value = loop = unittest.mock.Mock(spec=("run",
"add_idle"))
with _runstate.state_context(None):
_eventloop.add_idle("foo", "bar", baz="qux")
loop.add_idle.assert_called_once_with("foo", "bar", baz="qux")
def test_add_idle():
with pytest.raises(NotImplementedError):
eventloop.add_idle()