multi-instance test compatibility changes (#89)

marqo-ai · May 19, 2023 · cf4677b · cf4677b
1 parent 68538f5
commit cf4677b
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 2 deletions.
diff --git a/tests/v0_tests/test_bulk_search.py b/tests/v0_tests/test_bulk_search.py
@@ -47,6 +47,13 @@ def test_search_single(self):
             """
         }
         add_doc_res = self.client.index(self.index_name_1).add_documents([d1])
+
+        if self.IS_MULTI_INSTANCE:
+            self.warm_request(self.client.bulk_search, [{
+                "index": self.index_name_1,
+                "q": "title about some doc"
+            }])
+
         resp = self.client.bulk_search([{
             "index": self.index_name_1,
             "q": "title about some doc"
@@ -89,6 +96,13 @@ def test_search_highlights(self):
                 ({"showHighlights": False}, False),
                 ({"showHighlights": True}, True)
             ]:
+
+            if self.IS_MULTI_INSTANCE:
+                self.warm_request(self.client.bulk_search, [{**{
+                    "index": self.index_name_1,
+                    "q": "title about some doc"
+                }, **params}])
+
             resp = self.client.bulk_search([{**{
                 "index": self.index_name_1,
                 "q": "title about some doc"
@@ -112,6 +126,13 @@ def test_search_multi(self):
         res = self.client.index(self.index_name_1).add_documents([
             d1, d2
         ])
+
+        if self.IS_MULTI_INSTANCE:
+            self.warm_request(self.client.bulk_search, [{
+                "index": self.index_name_1,
+                "q": "this is a solid doc"
+            }])
+
         resp = self.client.bulk_search([{
             "index": self.index_name_1,
             "q": "this is a solid doc"
@@ -139,6 +160,12 @@ def test_select_lexical(self):
             d1, d2
         ])
 
+        if self.IS_MULTI_INSTANCE:
+            self.warm_request(self.client.bulk_search, [{
+                "index": self.index_name_1,
+                "q": "Examples of leadership"
+            }])
+
         # Ensure that vector search works
         resp = self.client.bulk_search([{
             "index": self.index_name_1,
@@ -204,6 +231,13 @@ def test_prefiltering(self):
             d1, d2
         ],auto_refresh=True)
 
+        if self.IS_MULTI_INSTANCE:
+            self.warm_request(self.client.bulk_search, [{
+                "index": self.index_name_1,
+                "q": "blah blah",
+                "filter": "(an_int:[0 TO 30] and an_int:2) AND abc-123:(some text)"
+            }])
+
         resp = self.client.bulk_search([{
             "index": self.index_name_1,
             "q": "blah blah",
@@ -237,6 +271,15 @@ def test_attributes_to_retrieve(self):
         atts = ["doc title", "an_int"]
         for search_method in [enums.SearchMethods.TENSOR,
                               enums.SearchMethods.LEXICAL]:
+
+            if self.IS_MULTI_INSTANCE:
+                self.warm_request(self.client.bulk_search,[{
+                    "index": self.index_name_1,
+                    "q": "blah blah",
+                    "searchMethod": search_method,
+                    "attributesToRetrieve": atts
+                }])
+
             resp = self.client.bulk_search([{
                 "index": self.index_name_1,
                 "q": "blah blah",

diff --git a/tests/v0_tests/test_score_modifier_search.py b/tests/v0_tests/test_score_modifier_search.py
@@ -55,6 +55,9 @@ def test_score_modifier_search_results(self):
                      }]
             }
 
+        if self.IS_MULTI_INSTANCE:
+            self.warm_request(self.client.index(self.index_name_1).search, q = self.query, score_modifiers=None, filter_string="filter:original")
+
         original_res = self.client.index(self.index_name_1).search(q = self.query, score_modifiers=None,
                                                                      filter_string="filter:original")
         original_score = original_res["hits"][0]["_score"]

diff --git a/tests/v0_tests/test_tensor_search.py b/tests/v0_tests/test_tensor_search.py
@@ -9,6 +9,7 @@
 import requests
 import random
 import math
+import time
 from tests.marqo_test import MarqoTestCase
 
 
@@ -137,6 +138,7 @@ def test_select_lexical(self):
 
         # Ensure that vector search works
         if self.IS_MULTI_INSTANCE:
+            time.sleep(5)
             self.warm_request(self.client.index(self.index_name_1).search,
             "Examples of leadership", search_method=enums.SearchMethods.TENSOR)
 
@@ -375,9 +377,16 @@ def test_special_characters(self):
             add_res = self.client.index(index_name=self.index_name_1).add_documents(
                 docs, auto_refresh=True, non_tensor_fields=[field_to_not_search]
             )
-            if self.IS_MULTI_INSTANCE:
-                self.warm_request(self.client.index(self.index_name_1).search, q='Blah')
+
             search_filter_field = f"filter{filter_char}me"
+            if self.IS_MULTI_INSTANCE:
+                self.warm_request(self.client.index(self.index_name_1).search, 
+                    q="Dog", 
+                    searchable_attributes=[field_to_search, field_to_not_search],
+                    attributes_to_retrieve=[field_to_not_search],
+                    filter_string=f'{search_filter_field}:Walrus'
+                )
+
             search1_res = self.client.index(index_name=self.index_name_1).search(
                 "Dog", searchable_attributes=[field_to_search, field_to_not_search],
                 attributes_to_retrieve=[field_to_not_search],