examples integration tests

Author: miguelgrinberg

examples/alias_migration.py

@@ -35,7 +35,7 @@
       will have index set to the concrete index whereas the class refers to the
       alias.
 """
-import os
+import asyncio
 from datetime import datetime
 from fnmatch import fnmatch
 
@@ -125,9 +125,9 @@ def migrate(move_data=True, update_alias=True):
         )
 
 
-if __name__ == "__main__":
+def main():
     # initiate the default connection to elasticsearch
-    connections.create_connection(hosts=[os.environ["ELASTICSEARCH_URL"]])
+    connections.create_connection(hosts=["http://localhost:9200"])
 
     # create the empty index
     setup()
@@ -143,3 +143,10 @@ def migrate(move_data=True, update_alias=True):
 
     # create new index
     migrate()
+
+    # close the connection
+    connections.get_connection().close()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

examples/async/parent_child.py

@@ -159,14 +159,14 @@ def search_answers(self):
         s = s.params(routing=self.meta.id)
         return s
 
-    def get_answers(self):
+    async def get_answers(self):
         """
         Get answers either from inner_hits already present or by searching
         elasticsearch.
         """
         if "inner_hits" in self.meta and "answer" in self.meta.inner_hits:
             return self.meta.inner_hits.answer.hits
-        return list(self.search_answers())
+        return [a async for a in self.search_answers()]
 
     async def save(self, **kwargs):
         self.question_answer = "question"
@@ -188,8 +188,7 @@ def _matches(cls, hit):
     def search(cls, **kwargs):
         return cls._index.search(**kwargs).exclude("term", question_answer="question")
 
-    @property
-    async def question(self):
+    async def get_question(self):
         # cache question in self.meta
         # any attributes set on self would be interpretted as fields
         if "question" not in self.meta:

examples/async/percolate.py

@@ -20,10 +20,10 @@
 
 from elasticsearch_dsl import (
     AsyncDocument,
+    AsyncSearch,
     Keyword,
     Percolator,
     Q,
-    Search,
     Text,
     async_connections,
 )
@@ -40,22 +40,22 @@ class BlogPost(AsyncDocument):
     class Index:
         name = "test-blogpost"
 
-    def add_tags(self):
+    async def add_tags(self):
         # run a percolation to automatically tag the blog post.
-        s = Search(index="test-percolator")
+        s = AsyncSearch(index="test-percolator")
         s = s.query(
             "percolate", field="query", index=self._get_index(), document=self.to_dict()
         )
 
         # collect all the tags from matched percolators
-        for percolator in s:
+        async for percolator in s:
             self.tags.extend(percolator.tags)
 
         # make sure tags are unique
         self.tags = list(set(self.tags))
 
     async def save(self, **kwargs):
-        self.add_tags()
+        await self.add_tags()
         return await super().save(**kwargs)
 
 
@@ -82,7 +82,7 @@ class Index:
 async def setup():
     # create the percolator index if it doesn't exist
     if not await PercolatorDoc._index.exists():
-        PercolatorDoc.init()
+        await PercolatorDoc.init()
 
     # register a percolation query looking for documents about python
     await PercolatorDoc(

examples/completion.py

@@ -26,6 +26,7 @@
 that does ascii folding.
 """
 
+import asyncio
 import os
 from itertools import permutations
 
@@ -72,7 +73,7 @@ class Index:
         settings = {"number_of_shards": 1, "number_of_replicas": 0}
 
 
-if __name__ == "__main__":
+def main():
     # initiate the default connection to elasticsearch
     connections.create_connection(hosts=[os.environ["ELASTICSEARCH_URL"]])
 
@@ -97,3 +98,10 @@ class Index:
         # print out all the options we got
         for option in response.suggest.auto_complete[0].options:
             print("%10s: %25s (%d)" % (text, option._source.name, option._score))
+
+    # close the connection
+    connections.get_connection().close()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

examples/composite_agg.py

@@ -15,6 +15,7 @@
 #  specific language governing permissions and limitations
 #  under the License.
 
+import asyncio
 import os
 
 from elasticsearch_dsl import A, Search, connections
@@ -36,15 +37,16 @@ def run_search(**kwargs):
 
     response = run_search()
     while response.aggregations.comp.buckets:
-        yield from response.aggregations.comp.buckets
+        for b in response.aggregations.comp.buckets:
+            yield b
         if "after_key" in response.aggregations.comp:
             after = response.aggregations.comp.after_key
         else:
             after = response.aggregations.comp.buckets[-1].key
         response = run_search(after=after)
 
 
-if __name__ == "__main__":
+def main():
     # initiate the default connection to elasticsearch
     connections.create_connection(hosts=[os.environ["ELASTICSEARCH_URL"]])
 
@@ -57,3 +59,10 @@ def run_search(**kwargs):
             "File %s has been modified %d times, first seen at %s."
             % (b.key.files, b.doc_count, b.first_seen.value_as_string)
         )
+
+    # close the connection
+    connections.get_connection().close()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

examples/parent_child.py

@@ -39,6 +39,7 @@
           particular parent
 
 """
+import asyncio
 import os
 from datetime import datetime
 
@@ -165,7 +166,7 @@ def get_answers(self):
         """
         if "inner_hits" in self.meta and "answer" in self.meta.inner_hits:
             return self.meta.inner_hits.answer.hits
-        return list(self.search_answers())
+        return [a for a in self.search_answers()]
 
     def save(self, **kwargs):
         self.question_answer = "question"
@@ -187,8 +188,7 @@ def _matches(cls, hit):
     def search(cls, **kwargs):
         return cls._index.search(**kwargs).exclude("term", question_answer="question")
 
-    @property
-    def question(self):
+    def get_question(self):
         # cache question in self.meta
         # any attributes set on self would be interpretted as fields
         if "question" not in self.meta:
@@ -209,7 +209,7 @@ def setup():
     index_template.save()
 
 
-if __name__ == "__main__":
+def main():
     # initiate the default connection to elasticsearch
     connections.create_connection(hosts=[os.environ["ELASTICSEARCH_URL"]])
 
@@ -244,3 +244,12 @@ def setup():
     )
     question.save()
     answer = question.add_answer(honza, "Just use `elasticsearch-py`!")
+
+    # close the connection
+    connections.get_connection().close()
+
+    return answer
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

examples/percolate.py

@@ -15,6 +15,7 @@
 #  specific language governing permissions and limitations
 #  under the License.
 
+import asyncio
 import os
 
 from elasticsearch_dsl import (
@@ -91,8 +92,15 @@ def setup():
     ).save(refresh=True)
 
 
-if __name__ == "__main__":
+def main():
     # initiate the default connection to elasticsearch
     connections.create_connection(hosts=[os.environ["ELASTICSEARCH_URL"]])
 
     setup()
+
+    # close the connection
+    connections.get_connection().close()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

tests/test_integration/test_examples/_async/__init__.py

@@ -0,0 +1,16 @@
+#  Licensed to Elasticsearch B.V. under one or more contributor
+#  license agreements. See the NOTICE file distributed with
+#  this work for additional information regarding copyright
+#  ownership. Elasticsearch B.V. licenses this file to you under
+#  the Apache License, Version 2.0 (the "License"); you may
+#  not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+# 	http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing,
+#  software distributed under the License is distributed on an
+#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+#  KIND, either express or implied.  See the License for the
+#  specific language governing permissions and limitations
+#  under the License.

tests/test_integration/test_examples/_async/alias_migration.py

@@ -0,0 +1 @@
+../../../../examples/async/alias_migration.py

tests/test_integration/test_examples/_async/completion.py

@@ -0,0 +1 @@
+../../../../examples/async/completion.py

tests/test_integration/test_examples/_async/composite_agg.py

@@ -0,0 +1 @@
+../../../../examples/async/composite_agg.py

tests/test_integration/test_examples/_async/parent_child.py

@@ -0,0 +1 @@
+../../../../examples/async/parent_child.py

tests/test_integration/test_examples/_async/percolate.py

@@ -0,0 +1 @@
+../../../../examples/async/percolate.py

tests/test_integration/test_examples/_async/test_alias_migration.py

@@ -0,0 +1,68 @@
+#  Licensed to Elasticsearch B.V. under one or more contributor
+#  license agreements. See the NOTICE file distributed with
+#  this work for additional information regarding copyright
+#  ownership. Elasticsearch B.V. licenses this file to you under
+#  the Apache License, Version 2.0 (the "License"); you may
+#  not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+# 	http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing,
+#  software distributed under the License is distributed on an
+#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+#  KIND, either express or implied.  See the License for the
+#  specific language governing permissions and limitations
+#  under the License.
+
+from . import alias_migration
+from .alias_migration import ALIAS, PATTERN, BlogPost, migrate
+
+
+async def test_alias_migration(async_write_client):
+    # create the index
+    await alias_migration.setup()
+
+    # verify that template, index, and alias has been set up
+    assert await async_write_client.indices.exists_template(name=ALIAS)
+    assert await async_write_client.indices.exists(index=PATTERN)
+    assert await async_write_client.indices.exists_alias(name=ALIAS)
+
+    indices = await async_write_client.indices.get(index=PATTERN)
+    assert len(indices) == 1
+    index_name, _ = indices.popitem()
+
+    # which means we can now save a document
+    with open(__file__) as f:
+        bp = BlogPost(
+            _id=0,
+            title="Hello World!",
+            tags=["testing", "dummy"],
+            content=f.read(),
+        )
+        await bp.save(refresh=True)
+
+    assert await BlogPost.search().count() == 1
+
+    # _matches work which means we get BlogPost instance
+    bp = (await BlogPost.search().execute())[0]
+    assert isinstance(bp, BlogPost)
+    assert not bp.is_published()
+    assert "0" == bp.meta.id
+
+    # create new index
+    await migrate()
+
+    indices = await async_write_client.indices.get(index=PATTERN)
+    assert 2 == len(indices)
+    alias = await async_write_client.indices.get(index=ALIAS)
+    assert 1 == len(alias)
+    assert index_name not in alias
+
+    # data has been moved properly
+    assert await BlogPost.search().count() == 1
+
+    # _matches work which means we get BlogPost instance
+    bp = (await BlogPost.search().execute())[0]
+    assert isinstance(bp, BlogPost)
+    assert "0" == bp.meta.id

tests/test_integration/test_examples/_async/test_completion.py

@@ -0,0 +1,34 @@
+#  Licensed to Elasticsearch B.V. under one or more contributor
+#  license agreements. See the NOTICE file distributed with
+#  this work for additional information regarding copyright
+#  ownership. Elasticsearch B.V. licenses this file to you under
+#  the Apache License, Version 2.0 (the "License"); you may
+#  not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+# 	http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing,
+#  software distributed under the License is distributed on an
+#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+#  KIND, either express or implied.  See the License for the
+#  specific language governing permissions and limitations
+#  under the License.
+
+
+from .completion import Person
+
+
+async def test_person_suggests_on_all_variants_of_name(async_write_client):
+    await Person.init(using=async_write_client)
+
+    await Person(name="Honza Král", popularity=42).save(refresh=True)
+
+    s = Person.search().suggest("t", "kra", completion={"field": "suggest"})
+    response = await s.execute()
+
+    opts = response.suggest.t[0].options
+
+    assert 1 == len(opts)
+    assert opts[0]._score == 42
+    assert opts[0]._source.name == "Honza Král"