add a unit test using some of the generated classes

miguelgrinberg · Sep 3, 2024 · 4cb5143 · 4cb5143
1 parent 6660e13
commit 4cb5143
Show file tree

Hide file tree

Showing 2 changed files with 154 additions and 2 deletions.
diff --git a/tests/_async/test_search.py b/tests/_async/test_search.py
@@ -21,7 +21,15 @@
 import pytest
 from pytest import raises
 
-from elasticsearch_dsl import AsyncEmptySearch, AsyncSearch, Document, Q, query
+from elasticsearch_dsl import (
+    AsyncEmptySearch,
+    AsyncSearch,
+    Document,
+    Q,
+    interfaces,
+    query,
+    wrappers,
+)
 from elasticsearch_dsl.exceptions import IllegalOperation
 
 
@@ -530,6 +538,74 @@ def test_reverse() -> None:
     assert d == s.to_dict()
 
 
+def test_code_generated_classes() -> None:
+    s = AsyncSearch()
+    s = (
+        s.query(query.Match("title", interfaces.MatchQuery(query="python")))
+        .query(~query.Match("title", interfaces.MatchQuery(query="ruby")))
+        .query(
+            query.Knn(
+                field="title",
+                query_vector=[1.0, 2.0, 3.0],
+                num_candidates=10,
+                k=3,
+                filter=query.Range("year", wrappers.Range(gt="2004")),
+            )
+        )
+        .filter(
+            query.Term("category", interfaces.TermQuery(value="meetup"))
+            | query.Term("category", interfaces.TermQuery(value="conference"))
+        )
+        .collapse("user_id")
+        .post_filter(query.Terms(tags=["prague", "czech"]))
+        .script_fields(more_attendees="doc['attendees'].value + 42")
+    )
+    assert {
+        "query": {
+            "bool": {
+                "filter": [
+                    {
+                        "bool": {
+                            "should": [
+                                {"term": {"category": {"value": "meetup"}}},
+                                {"term": {"category": {"value": "conference"}}},
+                            ]
+                        }
+                    }
+                ],
+                "must": [
+                    {"match": {"title": {"query": "python"}}},
+                    {
+                        "knn": {
+                            "field": "title",
+                            "filter": [
+                                {
+                                    "range": {
+                                        "year": {
+                                            "gt": "2004",
+                                        },
+                                    },
+                                },
+                            ],
+                            "k": 3,
+                            "num_candidates": 10,
+                            "query_vector": [
+                                1.0,
+                                2.0,
+                                3.0,
+                            ],
+                        },
+                    },
+                ],
+                "must_not": [{"match": {"title": {"query": "ruby"}}}],
+            }
+        },
+        "post_filter": {"terms": {"tags": ["prague", "czech"]}},
+        "collapse": {"field": "user_id"},
+        "script_fields": {"more_attendees": {"script": "doc['attendees'].value + 42"}},
+    } == s.to_dict()
+
+
 def test_from_dict_doesnt_need_query() -> None:
     s = AsyncSearch.from_dict({"size": 5})
 

diff --git a/tests/_sync/test_search.py b/tests/_sync/test_search.py
@@ -21,7 +21,15 @@
 import pytest
 from pytest import raises
 
-from elasticsearch_dsl import Document, EmptySearch, Q, Search, query
+from elasticsearch_dsl import (
+    Document,
+    EmptySearch,
+    Q,
+    Search,
+    interfaces,
+    query,
+    wrappers,
+)
 from elasticsearch_dsl.exceptions import IllegalOperation
 
 
@@ -530,6 +538,74 @@ def test_reverse() -> None:
     assert d == s.to_dict()
 
 
+def test_code_generated_classes() -> None:
+    s = Search()
+    s = (
+        s.query(query.Match("title", interfaces.MatchQuery(query="python")))
+        .query(~query.Match("title", interfaces.MatchQuery(query="ruby")))
+        .query(
+            query.Knn(
+                field="title",
+                query_vector=[1.0, 2.0, 3.0],
+                num_candidates=10,
+                k=3,
+                filter=query.Range("year", wrappers.Range(gt="2004")),
+            )
+        )
+        .filter(
+            query.Term("category", interfaces.TermQuery(value="meetup"))
+            | query.Term("category", interfaces.TermQuery(value="conference"))
+        )
+        .collapse("user_id")
+        .post_filter(query.Terms(tags=["prague", "czech"]))
+        .script_fields(more_attendees="doc['attendees'].value + 42")
+    )
+    assert {
+        "query": {
+            "bool": {
+                "filter": [
+                    {
+                        "bool": {
+                            "should": [
+                                {"term": {"category": {"value": "meetup"}}},
+                                {"term": {"category": {"value": "conference"}}},
+                            ]
+                        }
+                    }
+                ],
+                "must": [
+                    {"match": {"title": {"query": "python"}}},
+                    {
+                        "knn": {
+                            "field": "title",
+                            "filter": [
+                                {
+                                    "range": {
+                                        "year": {
+                                            "gt": "2004",
+                                        },
+                                    },
+                                },
+                            ],
+                            "k": 3,
+                            "num_candidates": 10,
+                            "query_vector": [
+                                1.0,
+                                2.0,
+                                3.0,
+                            ],
+                        },
+                    },
+                ],
+                "must_not": [{"match": {"title": {"query": "ruby"}}}],
+            }
+        },
+        "post_filter": {"terms": {"tags": ["prague", "czech"]}},
+        "collapse": {"field": "user_id"},
+        "script_fields": {"more_attendees": {"script": "doc['attendees'].value + 42"}},
+    } == s.to_dict()
+
+
 def test_from_dict_doesnt_need_query() -> None:
     s = Search.from_dict({"size": 5})