Refactor unit test.

Author: WaVEV

tests/queries_/models.py

@@ -65,10 +65,15 @@ class Writer(EmbeddedModel):
     name = models.CharField(max_length=10)
 
 
+class Location(EmbeddedModel):
+    type = models.CharField(default="Point", max_length=50)
+    coordinates = ArrayField(models.FloatField(), max_size=2)
+
+
 class Article(models.Model):
     headline = models.CharField(max_length=100)
     number = models.IntegerField()
     body = models.TextField()
-    location = models.JSONField(null=True)
+    location = EmbeddedModelField(Location, null=True)
     plot_embedding = ArrayField(models.FloatField(), size=3, null=True)
     writer = EmbeddedModelField(Writer, null=True)

tests/queries_/test_search.py

@@ -28,7 +28,18 @@
     SearchWildcard,
 )
 
-from .models import Article, Writer
+from .models import Article, Location, Writer
+
+
+def wait_until_index_ready(collection, index_name, timeout: float = 30, interval: float = 0.5):
+    start = monotonic()
+    while monotonic() - start < timeout:
+        indexes = list(collection.list_search_indexes())
+        for idx in indexes:
+            if idx["name"] == index_name and idx["status"] == "READY":
+                return True
+        sleep(interval)
+    raise TimeoutError(f"Index {index_name} not ready after {timeout} seconds")
 
 
 def _delayed_assertion(timeout: float = 120, interval: float = 0.5):
@@ -76,6 +87,7 @@ def create_search_index(cls, model, index_name, definition, type="search"):
         collection = cls._get_collection(model)
         idx = SearchIndexModel(definition=definition, name=index_name, type=type)
         collection.create_search_index(idx)
+        wait_until_index_ready(collection, index_name)
 
         def drop_index():
             collection.drop_search_index(index_name)
@@ -479,10 +491,13 @@ def setUpClass(cls):
 
     def setUp(self):
         self.article = Article.objects.create(
-            headline="any", number=1, body="", location={"type": "Point", "coordinates": [40, 5]}
+            headline="any", number=1, body="", location=Location(type="Point", coordinates=[40, 5])
         )
         Article.objects.create(
-            headline="any", number=2, body="", location={"type": "Point", "coordinates": [400, 50]}
+            headline="any",
+            number=2,
+            body="",
+            location=Location(type="Point", coordinates=[400, 50]),
         )
 
     def test_search_geo_shape(self):
@@ -524,10 +539,13 @@ def setUpClass(cls):
 
     def setUp(self):
         self.article = Article.objects.create(
-            headline="geo", number=2, body="", location={"type": "Point", "coordinates": [40, 5]}
+            headline="geo", number=2, body="", location=Location(type="Point", coordinates=[40, 5])
         )
         Article.objects.create(
-            headline="geo2", number=3, body="", location={"type": "Point", "coordinates": [-40, -5]}
+            headline="geo2",
+            number=3,
+            body="",
+            location=Location(type="Point", coordinates=[-40, -5]),
         )
 
     def test_search_geo_within(self):
@@ -598,7 +616,6 @@ def test_search_more_like_this(self):
             {"headline": self.article1.headline, "body": self.article1.body},
             {"headline": self.article2.headline, "body": self.article2.body},
         ]
-        like_docs = [{"body": "NASA launches new satellite to explore the galaxy"}]
         qs = Article.objects.annotate(score=SearchMoreLikeThis(documents=like_docs)).order_by(
             "score"
         )