def test_found_missing_deferred(runstate):
def key_pb(key):
mock_key = mock.Mock(spec=("SerializeToString",))
mock_key.SerializeToString.return_value = key
return mock_key
runstate.eventloop = mock.Mock(spec=("add_idle", "run"))
future1, future2, future3 = (tasklets.Future() for _ in range(3))
batch = _api._LookupBatch({})
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 = runstate.batches[_api._LookupBatch][()]
assert next_batch.todo == {"baz": [future3]}
assert runstate.eventloop.add_idle.call_count == 1
def test_deferred(runstate):
def key_pb(key):
mock_key = mock.Mock(spec=("SerializeToString",))
mock_key.SerializeToString.return_value = key
return mock_key
runstate.eventloop = mock.Mock(spec=("add_idle", "run"))
future1, future2, future3 = (tasklets.Future() for _ in range(3))
batch = _api._LookupBatch({})
batch.todo.update({"foo": [future1, future2], "bar": [future3]})
response = mock.Mock(
missing=[],
found=[],
deferred=[key_pb("foo"), key_pb("bar")],
spec=("found", "missing", "deferred"),
)
rpc = tasklets.Future()
rpc.set_result(response)
batch.lookup_callback(rpc)
assert future1.running()
assert future2.running()
assert future3.running()
next_batch = runstate.batches[_api._LookupBatch][()]
assert next_batch.todo == batch.todo and next_batch is not batch
assert runstate.eventloop.add_idle.call_count == 1
def test_missing():
def key_pb(key):
mock_key = mock.Mock(spec=("SerializeToString",))
mock_key.SerializeToString.return_value = key
return mock_key
future1, future2, future3 = (tasklets.Future() for _ in range(3))
batch = _api._LookupBatch({})
batch.todo.update({"foo": [future1, future2], "bar": [future3]})
entity1 = mock.Mock(key=key_pb("foo"), spec=("key",))
entity2 = mock.Mock(key=key_pb("bar"), spec=("key",))
response = mock.Mock(
missing=[
mock.Mock(entity=entity1, spec=("entity",)),
mock.Mock(entity=entity2, spec=("entity",)),
],
found=[],
deferred=[],
spec=("found", "missing", "deferred"),
)
rpc = tasklets.Future()
rpc.set_result(response)
batch.lookup_callback(rpc)
assert future1.result() is _api._NOT_FOUND
assert future2.result() is _api._NOT_FOUND
assert future3.result() is _api._NOT_FOUND
def test_idle_callback(_datastore_lookup, entity_pb2, context):
class MockKey:
def __init__(self, key=None):
self.key = key
def ParseFromString(self, key):
self.key = key
rpc = tasklets.Future("_datastore_lookup")
_datastore_lookup.return_value = rpc
entity_pb2.Key = MockKey
eventloop = mock.Mock(spec=("queue_rpc", "run"))
with context.new(eventloop=eventloop).use() as context:
batch = _api._LookupBatch(_options.ReadOptions())
batch.lookup_callback = mock.Mock()
batch.todo.update({"foo": ["one", "two"], "bar": ["three"]})
batch.idle_callback()
called_with = _datastore_lookup.call_args[0]
called_with_keys = set(
(mock_key.key for mock_key in called_with[0])
)
assert called_with_keys == set(["foo", "bar"])
called_with_options = called_with[1]
assert called_with_options == datastore_pb2.ReadOptions()
rpc.set_result(None)
batch.lookup_callback.assert_called_once_with(rpc)
def test_lookup_callback_exception():
future1, future2, future3 = (tasklets.Future() for _ in range(3))
batch = _api._LookupBatch({})
batch.todo.update({"foo": [future1, future2], "bar": [future3]})
error = Exception("Spurious error.")
rpc = tasklets.Future()
rpc.set_exception(error)
batch.lookup_callback(rpc)
assert future1.exception() is error
assert future2.exception() is error
def test_idle_callback(_datastore_lookup, entity_pb2, runstate):
class MockKey:
def __init__(self, key=None):
self.key = key
def ParseFromString(self, key):
self.key = key
entity_pb2.Key = MockKey
runstate.eventloop = mock.Mock(spec=("queue_rpc", "run"))
batch = _api._LookupBatch({})
batch.todo.update({"foo": ["one", "two"], "bar": ["three"]})
batch.idle_callback()
called_with = _datastore_lookup.call_args[0]
called_with_keys = set((mock_key.key for mock_key in called_with[0]))
assert called_with_keys == set(["foo", "bar"])
called_with_options = called_with[1]
assert called_with_options == datastore_pb2.ReadOptions()
rpc = _datastore_lookup.return_value
runstate.eventloop.queue_rpc.assert_called_once_with(
rpc, batch.lookup_callback
)