added tests

Author: daniel-sanche

packages/google-cloud-firestore/tests/system/pipeline_e2e/search.yaml

@@ -22,3 +22,133 @@ tests:
                 args:
                 - stringValue: "technology"
                 name: document_matches
+  - description: search_stage_full_options
+    pipeline:
+      - Collection: books
+      - Search:
+          - SearchOptions:
+              query:
+                DocumentMatches:
+                  - Constant: "science"
+              limit: 5
+              retrieval_depth: 10
+              offset: 2
+              query_enhancement: disabled
+              language_code: en
+    assert_proto:
+      pipeline:
+        stages:
+        - args:
+          - referenceValue: /books
+          name: collection
+        - args: []
+          name: search
+          options:
+            limit:
+              integerValue: '5'
+            retrieval_depth:
+              integerValue: '10'
+            offset:
+              integerValue: '2'
+            query_enhancement:
+              stringValue: "disabled"
+            language_code:
+              stringValue: "en"
+            query:
+              functionValue:
+                args:
+                - stringValue: "science"
+                name: document_matches
+  - description: search_stage_with_sort_and_select
+    pipeline:
+      - Collection: books
+      - Search:
+          - SearchOptions:
+              query: "science"
+              sort:
+                Ordering:
+                  - Field: rating
+                  - DESCENDING
+              select:
+                - title
+              add_fields:
+                - AliasedExpression:
+                  - FunctionExpression.string_concat:
+                    - Field: title
+                    - Constant: " - Selected"
+                  - "title_selected"
+    assert_proto:
+      pipeline:
+        stages:
+        - args:
+          - referenceValue: /books
+          name: collection
+        - args: []
+          name: search
+          options:
+            query:
+              functionValue:
+                args:
+                - stringValue: "science"
+                name: document_matches
+            sort:
+              arrayValue:
+                values:
+                - mapValue:
+                    fields:
+                      direction:
+                        stringValue: descending
+                      expression:
+                        fieldReferenceValue: rating
+            select:
+              mapValue:
+                fields:
+                  title:
+                    fieldReferenceValue: title
+            add_fields:
+              mapValue:
+                fields:
+                  title_selected:
+                    functionValue:
+                      args:
+                      - fieldReferenceValue: title
+                      - stringValue: " - Selected"
+                      name: string_concat
+  - description: expression_between
+    pipeline:
+      - Collection: books
+      - Where:
+          - FunctionExpression.between:
+              - Field: published
+              - Constant: 1950
+              - Constant: 1970
+      - Select:
+        - title
+      - Sort:
+        - Ordering:
+          - Field: title
+          - ASCENDING
+    assert_results:
+      - title: Dune
+      - title: One Hundred Years of Solitude
+      - title: The Lord of the Rings
+      - title: To Kill a Mockingbird
+  - description: expression_geo_distance
+    pipeline:
+      - Collection: geopoints
+      - Search:
+          - SearchOptions:
+              query:
+                FunctionExpression.less_than:
+                  - FunctionExpression.geo_distance:
+                      - Field: location
+                      - Constant: GEOPOINT(37.0,-122.0)
+                  - Constant: 150000.0
+      - Select:
+        - name
+      - Sort:
+        - Ordering:
+          - Field: name
+          - ASCENDING
+    assert_results:
+      - name: SF

packages/google-cloud-firestore/tests/unit/v1/test_pipeline_stages.py

@@ -24,6 +24,7 @@
     Constant,
     Field,
     Ordering,
+    DocumentMatches,
 )
 from google.cloud.firestore_v1.types.document import Value
 from google.cloud.firestore_v1.vector import Vector
@@ -778,6 +779,68 @@ def test_to_pb_percent_mode(self):
         assert len(result_percent.options) == 0
 
 
+class TestSearch:
+    def test_search_defaults(self):
+        options = stages.SearchOptions(query="technology")
+        assert options.query.name == "document_matches"
+        assert options.limit is None
+        assert options.retrieval_depth is None
+        assert options.sort is None
+        assert options.add_fields is None
+        assert options.select is None
+        assert options.offset is None
+        assert options.query_enhancement is None
+        assert options.language_code is None
+
+        stage = stages.Search(options)
+        pb_opts = stage._pb_options()
+        assert "query" in pb_opts
+        assert "limit" not in pb_opts
+        assert "retrieval_depth" not in pb_opts
+
+    def test_search_full_options(self):
+        options = stages.SearchOptions(
+            query=DocumentMatches("tech"),
+            limit=10,
+            retrieval_depth=2,
+            sort=Ordering("score", Ordering.Direction.DESCENDING),
+            add_fields=[Field("extra")],
+            select=[Field("name")],
+            offset=5,
+            query_enhancement="disabled",
+            language_code="en",
+        )
+        assert options.limit == 10
+        assert options.retrieval_depth == 2
+        assert len(options.sort) == 1
+        assert options.offset == 5
+        assert options.query_enhancement == stages.QueryEnhancement.DISABLED
+        assert options.language_code == "en"
+
+        stage = stages.Search(options)
+        pb_opts = stage._pb_options()
+
+        assert pb_opts["limit"].integer_value == 10
+        assert pb_opts["retrieval_depth"].integer_value == 2
+        assert len(pb_opts["sort"].array_value.values) == 1
+        assert pb_opts["offset"].integer_value == 5
+        assert pb_opts["query_enhancement"].string_value == "disabled"
+        assert pb_opts["language_code"].string_value == "en"
+        assert "query" in pb_opts
+
+    def test_search_string_query_wrapping(self):
+        options = stages.SearchOptions(query="science")
+        assert options.query.name == "document_matches"
+
+    def test_search_query_enhancement_enum(self):
+        options = stages.SearchOptions(query="q", query_enhancement=stages.QueryEnhancement.REQUIRED)
+        assert options.query_enhancement == stages.QueryEnhancement.REQUIRED
+        
+        stage = stages.Search(options)
+        pb_opts = stage._pb_options()
+        assert pb_opts["query_enhancement"].string_value == "required"
+
+
 class TestSelect:
     def _make_one(self, *args, **kwargs):
         return stages.Select(*args, **kwargs)