fix for search method

ionelmc · Jul 6, 2023 · 61c8961 · 61c8961
1 parent 3d8eafa
commit 61c8961
Showing 1 changed file with 13 additions and 5 deletions.
diff --git a/src/pytest_benchmark/storage/elasticsearch.py b/src/pytest_benchmark/storage/elasticsearch.py
@@ -13,9 +13,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]")
 
@@ -73,7 +73,7 @@ def query(self):
                 }
             }
         }
-        result = self._es.search(index=self._es_index, doc_type=self._es_doctype, body=body)
+        result = self._search_call(body=body)
         return sorted([record["key"] for record in result["aggregations"]["benchmark_ids"]["buckets"]])
 
     def load(self, id_prefix=None):
@@ -116,7 +116,9 @@ def _search(self, project, id_prefix=None):
                 }
             }
 
-        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):
@@ -178,13 +180,19 @@ def save(self, output_json, save):
         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": {