diff --git a/elasticmock/fake_elasticsearch.py b/elasticmock/fake_elasticsearch.py index 284e2e6..5a556d5 100644 --- a/elasticmock/fake_elasticsearch.py +++ b/elasticmock/fake_elasticsearch.py @@ -9,7 +9,6 @@ from elasticmock.utilities import get_random_id, get_random_scroll_id - PY3 = sys.version_info[0] == 3 if PY3: unicode = str @@ -185,6 +184,20 @@ def search(self, index=None, doc_type=None, body=None, params=None): match['_score'] = 1.0 hits.append(match) + # build aggregations + if body is not None and 'aggs' in body: + aggregations = {} + + for aggregation, definition in body['aggs'].items(): + aggregations[aggregation] = { + "doc_count_error_upper_bound": 0, + "sum_other_doc_count": 0, + "buckets": [] + } + + if aggregations: + result['aggregations'] = aggregations + if 'scroll' in params: result['_scroll_id'] = str(get_random_scroll_id()) params['size'] = int(params.get('size') if 'size' in params else 10) @@ -198,6 +211,7 @@ def search(self, index=None, doc_type=None, body=None, params=None): hits = hits[params.get('from'):params.get('from') + params.get('size')] result['hits']['hits'] = hits + return result @query_params('scroll') diff --git a/tests/test_elasticmock.py b/tests/test_elasticmock.py index 7840b15..27deaac 100644 --- a/tests/test_elasticmock.py +++ b/tests/test_elasticmock.py @@ -241,6 +241,14 @@ def test_doc_type_can_be_list(self): result = self.es.search(doc_type=doc_types[:2]) self.assertEqual(count_per_doc_type * 2, result.get('hits').get('total')) + def test_usage_of_aggregations(self): + self.es.index(index='index', doc_type='document', body={'genre': 'rock'}) + + body = {"aggs": {"genres": {"terms": {"field": "genre"}}}} + result = self.es.search(index='index', body=body) + + self.assertTrue('aggregations' in result) + def test_search_with_scroll_param(self): for _ in range(100): self.es.index(index='groups', doc_type='groups', body={'budget': 1000})