fix for search method

ionelmc · Dec 16, 2023 · cbeb977 · cbeb977
1 parent 8093bbc
commit cbeb977
Showing 1 changed file with 16 additions and 7 deletions.
diff --git a/src/pytest_benchmark/storage/elasticsearch.py b/src/pytest_benchmark/storage/elasticsearch.py
@@ -11,9 +11,9 @@
     import elasticsearch
     from elasticsearch.serializer import JSONSerializer
     api_version = elasticsearch.__version__
-    NO_DOC_TYPE_ARG_TO_INDEX_FUNC = False
+    NO_DOC_TYPE_ARG = False
     if api_version[0] > 7:
-        NO_DOC_TYPE_ARG_TO_INDEX_FUNC = True
+        NO_DOC_TYPE_ARG = True
 except ImportError:
     raise ImportError("Please install elasticsearch or pytest-benchmark[elasticsearch]")
 
@@ -61,8 +61,8 @@ def query(self):
         Returns sorted records names (ids) that corresponds with project.
         """
         body = {'size': 0, 'aggs': {'benchmark_ids': {'terms': {'field': 'benchmark_id'}}}}
-        result = self._es.search(index=self._es_index, doc_type=self._es_doctype, body=body)
-        return sorted([record['key'] for record in result['aggregations']['benchmark_ids']['buckets']])
+        result = self._search_call(body=body)
+        return sorted([record["key"] for record in result["aggregations"]["benchmark_ids"]["buckets"]])
 
     def load(self, id_prefix=None):
         """
@@ -86,7 +86,9 @@ def _search(self, project, id_prefix=None):
         if id_prefix:
             body['query']['bool']['must'] = {'prefix': {'_id': id_prefix}}
 
-        return self._es.search(index=self._es_index, doc_type=self._es_doctype, body=body)
+        result = self._search_call(body=body)
+
+        return result
 
     @staticmethod
     def _benchmark_from_es_record(source_es_record):
@@ -145,15 +147,22 @@ def save(self, output_json, save):
             )
         # hide user's credentials before logging
         masked_hosts = _mask_hosts(self._es_hosts)
-        self.logger.info(f'Saved benchmark data to {masked_hosts} to index {self._es_index} as doctype {self._es_doctype}')
+        self.logger.info("Saved benchmark data to %s to index %s as doctype %s" % (
+            masked_hosts, self._es_index, self._es_doctype))
 
-    if NO_DOC_TYPE_ARG_TO_INDEX_FUNC:
+    if NO_DOC_TYPE_ARG:
         def _index_call(self, body, id):
             return self._es.index(index=self._es_index, body=body, id=id)
+
+        def _search_call(self, body):
+            return self._es.search(index=self._es_index, body=body)
     else:
         def _index_call(self, body, id):
             return self._es.index(index=self._es_index, doc_type=self._es_doctype, body=body, id=id)
 
+        def _search_call(self, body):
+            return self._es.search(index=self._es_index, doc_type=self._es_doctype, body=body)
+
     def _create_index(self):
         mapping = {
             'mappings': {