Skip to content

Commit cbeb977

Browse files
committed
fix for search method
1 parent 8093bbc commit cbeb977

File tree

1 file changed

+16
-7
lines changed

1 file changed

+16
-7
lines changed

src/pytest_benchmark/storage/elasticsearch.py

Lines changed: 16 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,9 @@
1111
import elasticsearch
1212
from elasticsearch.serializer import JSONSerializer
1313
api_version = elasticsearch.__version__
14-
NO_DOC_TYPE_ARG_TO_INDEX_FUNC = False
14+
NO_DOC_TYPE_ARG = False
1515
if api_version[0] > 7:
16-
NO_DOC_TYPE_ARG_TO_INDEX_FUNC = True
16+
NO_DOC_TYPE_ARG = True
1717
except ImportError:
1818
raise ImportError("Please install elasticsearch or pytest-benchmark[elasticsearch]")
1919

@@ -61,8 +61,8 @@ def query(self):
6161
Returns sorted records names (ids) that corresponds with project.
6262
"""
6363
body = {'size': 0, 'aggs': {'benchmark_ids': {'terms': {'field': 'benchmark_id'}}}}
64-
result = self._es.search(index=self._es_index, doc_type=self._es_doctype, body=body)
65-
return sorted([record['key'] for record in result['aggregations']['benchmark_ids']['buckets']])
64+
result = self._search_call(body=body)
65+
return sorted([record["key"] for record in result["aggregations"]["benchmark_ids"]["buckets"]])
6666

6767
def load(self, id_prefix=None):
6868
"""
@@ -86,7 +86,9 @@ def _search(self, project, id_prefix=None):
8686
if id_prefix:
8787
body['query']['bool']['must'] = {'prefix': {'_id': id_prefix}}
8888

89-
return self._es.search(index=self._es_index, doc_type=self._es_doctype, body=body)
89+
result = self._search_call(body=body)
90+
91+
return result
9092

9193
@staticmethod
9294
def _benchmark_from_es_record(source_es_record):
@@ -145,15 +147,22 @@ def save(self, output_json, save):
145147
)
146148
# hide user's credentials before logging
147149
masked_hosts = _mask_hosts(self._es_hosts)
148-
self.logger.info(f'Saved benchmark data to {masked_hosts} to index {self._es_index} as doctype {self._es_doctype}')
150+
self.logger.info("Saved benchmark data to %s to index %s as doctype %s" % (
151+
masked_hosts, self._es_index, self._es_doctype))
149152

150-
if NO_DOC_TYPE_ARG_TO_INDEX_FUNC:
153+
if NO_DOC_TYPE_ARG:
151154
def _index_call(self, body, id):
152155
return self._es.index(index=self._es_index, body=body, id=id)
156+
157+
def _search_call(self, body):
158+
return self._es.search(index=self._es_index, body=body)
153159
else:
154160
def _index_call(self, body, id):
155161
return self._es.index(index=self._es_index, doc_type=self._es_doctype, body=body, id=id)
156162

163+
def _search_call(self, body):
164+
return self._es.search(index=self._es_index, doc_type=self._es_doctype, body=body)
165+
157166
def _create_index(self):
158167
mapping = {
159168
'mappings': {

0 commit comments

Comments
 (0)