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 )