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