def test_query_iterator(self):
self._mock_datastore.run_query.side_effect = (
self.permanent_retriable_datastore_failure)
query_iterator = helper.QueryIterator("project", None, self._query,
self._mock_datastore)
self.assertRaises(RPCError, iter(query_iterator).next)
self.assertEqual(6, len(self._mock_datastore.run_query.call_args_list))
def test_query_iterator_with_non_retriable_failures(self):
self._mock_datastore.run_query.side_effect = (
self.non_retriable_datastore_failure)
query_iterator = helper.QueryIterator("project", None, self._query,
self._mock_datastore)
self.assertRaises(tuple(map(type, self._non_retriable_errors)),
iter(query_iterator).next)
self.assertEqual(1, len(self._mock_datastore.run_query.call_args_list))
def test_query_iterator_with_transient_failures(self):
self._mock_datastore.run_query.side_effect = (
self.transient_retriable_datastore_failure)
query_iterator = helper.QueryIterator("project", None, self._query,
self._mock_datastore)
fail_count = 5
self._transient_fail_count = fail_count
for _ in query_iterator:
pass
self.assertEqual(fail_count + 1,
len(self._mock_datastore.run_query.call_args_list))
def check_query_iterator(self, num_entities, batch_size, query):
"""A helper method to test the QueryIterator.
Args:
num_entities: number of entities contained in the fake datastore.
batch_size: the number of entities returned by fake datastore in one req.
query: the query to be executed
"""
entities = fake_datastore.create_entities(num_entities)
self._mock_datastore.run_query.side_effect = \
fake_datastore.create_run_query(entities, batch_size)
query_iterator = helper.QueryIterator("project", None, self._query,
self._mock_datastore)
i = 0
for entity in query_iterator:
self.assertEqual(entity, entities[i].entity)
i += 1
limit = query.limit.value if query.HasField('limit') else sys.maxsize
self.assertEqual(i, min(num_entities, limit))
