chore(tests): re-enable pipeline system tests on kokoro

google/cloud/firestore_v1/async_pipeline.py

@@ -68,7 +68,6 @@ async def execute(
         transaction: "AsyncTransaction" | None = None,
         read_time: datetime.datetime | None = None,
         explain_options: PipelineExplainOptions | None = None,
-        index_mode: str | None = None,
         additional_options: dict[str, Value | Constant] = {},
     ) -> PipelineSnapshot[PipelineResult]:
         """
@@ -87,10 +86,8 @@ async def execute(
             explain_options (Optional[:class:`~google.cloud.firestore_v1.query_profile.PipelineExplainOptions`]):
                 Options to enable query profiling for this query. When set,
                 explain_metrics will be available on the returned list.
-            index_mode (Optional[str]): Configures the pipeline to require a certain type of indexes to be present.
-                Firestore will reject the request if there is not appropiate indexes to serve the query.
             additional_options (Optional[dict[str, Value | Constant]]): Additional options to pass to the query.
-                These options will take precedence over method argument if there is a conflict (e.g. explain_options, index_mode)
+                These options will take precedence over method argument if there is a conflict (e.g. explain_options)
         """
         kwargs = {k: v for k, v in locals().items() if k != "self"}
         stream = AsyncPipelineStream(PipelineResult, self, **kwargs)
@@ -103,7 +100,6 @@ def stream(
         read_time: datetime.datetime | None = None,
         transaction: "AsyncTransaction" | None = None,
         explain_options: PipelineExplainOptions | None = None,
-        index_mode: str | None = None,
         additional_options: dict[str, Value | Constant] = {},
     ) -> AsyncPipelineStream[PipelineResult]:
         """
@@ -122,10 +118,8 @@ def stream(
             explain_options (Optional[:class:`~google.cloud.firestore_v1.query_profile.PipelineExplainOptions`]):
                 Options to enable query profiling for this query. When set,
                 explain_metrics will be available on the returned generator.
-            index_mode (Optional[str]): Configures the pipeline to require a certain type of indexes to be present.
-                Firestore will reject the request if there is not appropiate indexes to serve the query.
             additional_options (Optional[dict[str, Value | Constant]]): Additional options to pass to the query.
-                These options will take precedence over method argument if there is a conflict (e.g. explain_options, index_mode)
+                These options will take precedence over method argument if there is a conflict (e.g. explain_options)
         """
         kwargs = {k: v for k, v in locals().items() if k != "self"}
         return AsyncPipelineStream(PipelineResult, self, **kwargs)

google/cloud/firestore_v1/pipeline.py

@@ -65,7 +65,6 @@ def execute(
         transaction: "Transaction" | None = None,
         read_time: datetime.datetime | None = None,
         explain_options: PipelineExplainOptions | None = None,
-        index_mode: str | None = None,
         additional_options: dict[str, Value | Constant] = {},
     ) -> PipelineSnapshot[PipelineResult]:
         """
@@ -84,10 +83,8 @@ def execute(
             explain_options (Optional[:class:`~google.cloud.firestore_v1.query_profile.PipelineExplainOptions`]):
                 Options to enable query profiling for this query. When set,
                 explain_metrics will be available on the returned list.
-            index_mode (Optional[str]): Configures the pipeline to require a certain type of indexes to be present.
-                Firestore will reject the request if there is not appropiate indexes to serve the query.
             additional_options (Optional[dict[str, Value | Constant]]): Additional options to pass to the query.
-                These options will take precedence over method argument if there is a conflict (e.g. explain_options, index_mode)
+                These options will take precedence over method argument if there is a conflict (e.g. explain_options)
         """
         kwargs = {k: v for k, v in locals().items() if k != "self"}
         stream = PipelineStream(PipelineResult, self, **kwargs)
@@ -100,7 +97,6 @@ def stream(
         transaction: "Transaction" | None = None,
         read_time: datetime.datetime | None = None,
         explain_options: PipelineExplainOptions | None = None,
-        index_mode: str | None = None,
         additional_options: dict[str, Value | Constant] = {},
     ) -> PipelineStream[PipelineResult]:
         """
@@ -119,10 +115,8 @@ def stream(
             explain_options (Optional[:class:`~google.cloud.firestore_v1.query_profile.PipelineExplainOptions`]):
                 Options to enable query profiling for this query. When set,
                 explain_metrics will be available on the returned generator.
-            index_mode (Optional[str]): Configures the pipeline to require a certain type of indexes to be present.
-                Firestore will reject the request if there is not appropiate indexes to serve the query.
             additional_options (Optional[dict[str, Value | Constant]]): Additional options to pass to the query.
-                These options will take precedence over method argument if there is a conflict (e.g. explain_options, index_mode)
+                These options will take precedence over method argument if there is a conflict (e.g. explain_options)
         """
         kwargs = {k: v for k, v in locals().items() if k != "self"}
         return PipelineStream(PipelineResult, self, **kwargs)

google/cloud/firestore_v1/pipeline_result.py

@@ -178,7 +178,6 @@ def __init__(
         transaction: Transaction | AsyncTransaction | None,
         read_time: datetime.datetime | None,
         explain_options: PipelineExplainOptions | None,
-        index_mode: str | None,
         additional_options: dict[str, Constant | Value],
     ):
         # public
@@ -192,7 +191,6 @@ def __init__(
         self._explain_stats: ExplainStats | None = None
         self._explain_options: PipelineExplainOptions | None = explain_options
         self._return_type = return_type
-        self._index_mode = index_mode
         self._additonal_options = {
             k: v if isinstance(v, Value) else v._to_pb()
             for k, v in additional_options.items()
@@ -226,8 +224,6 @@ def _build_request(self) -> ExecutePipelineRequest:
         options = {}
         if self._explain_options:
             options["explain_options"] = self._explain_options._to_value()
-        if self._index_mode:
-            options["index_mode"] = Value(string_value=self._index_mode)
         if self._additonal_options:
             options.update(self._additonal_options)
         request = ExecutePipelineRequest(

tests/system/test__helpers.py

@@ -15,10 +15,8 @@
 EMULATOR_CREDS = EmulatorCreds()
 FIRESTORE_EMULATOR = os.environ.get(_FIRESTORE_EMULATOR_HOST) is not None
 FIRESTORE_OTHER_DB = os.environ.get("SYSTEM_TESTS_DATABASE", "system-tests-named-db")
-FIRESTORE_ENTERPRISE_DB = os.environ.get("ENTERPRISE_DATABASE", "enterprise-db")
+FIRESTORE_ENTERPRISE_DB = os.environ.get("ENTERPRISE_DATABASE", "enterprise-db-native")
 
 # run all tests against default database, and a named database
 TEST_DATABASES = [None, FIRESTORE_OTHER_DB]
 TEST_DATABASES_W_ENTERPRISE = TEST_DATABASES + [FIRESTORE_ENTERPRISE_DB]
-# TODO remove when kokoro fully supports enterprise mode/pipelines
-IS_KOKORO_TEST = os.getenv("KOKORO_JOB_NAME") is not None

tests/system/test_pipeline_acceptance.py

@@ -33,16 +33,10 @@
 
 from google.cloud.firestore import Client, AsyncClient
 
-from test__helpers import FIRESTORE_ENTERPRISE_DB, IS_KOKORO_TEST
+from test__helpers import FIRESTORE_ENTERPRISE_DB
 
 FIRESTORE_PROJECT = os.environ.get("GCLOUD_PROJECT")
 
-# TODO: enable kokoro tests when internal test project is whitelisted
-pytestmark = pytest.mark.skipif(
-    condition=IS_KOKORO_TEST,
-    reason="Pipeline tests are currently not supported by kokoro",
-)
-
 test_dir_name = os.path.dirname(__file__)
 
 id_format = (

tests/system/test_system.py

@@ -46,7 +46,6 @@
     ENTERPRISE_MODE_ERROR,
     TEST_DATABASES,
     TEST_DATABASES_W_ENTERPRISE,
-    IS_KOKORO_TEST,
     FIRESTORE_ENTERPRISE_DB,
 )
 
@@ -67,12 +66,6 @@ def _get_credentials_and_project():
 
 @pytest.fixture(scope="session")
 def database(request):
-    from test__helpers import FIRESTORE_ENTERPRISE_DB
-
-    # enterprise mode currently does not support RunQuery calls in prod on kokoro test project
-    # TODO: remove skip when kokoro test project supports full enterprise mode
-    if request.param == FIRESTORE_ENTERPRISE_DB and IS_KOKORO_TEST:
-        pytest.skip("enterprise mode does not support RunQuery on kokoro")
     return request.param
 
 
@@ -101,11 +94,6 @@ def verify_pipeline(query):
     """
     from google.cloud.firestore_v1.base_aggregation import BaseAggregationQuery
 
-    # return early on kokoro. Test project doesn't currently support pipelines
-    # TODO: enable pipeline verification when kokoro test project is whitelisted
-    if IS_KOKORO_TEST:
-        pytest.skip("skipping pipeline verification on kokoro")
-
     def _clean_results(results):
         if isinstance(results, dict):
             return {k: _clean_results(v) for k, v in results.items()}
@@ -1771,22 +1759,6 @@ def test_pipeline_explain_options_using_additional_options(
     assert "Execution:" in text_stats
 
 
-@pytest.mark.skipif(
-    FIRESTORE_EMULATOR, reason="Query profile not supported in emulator."
-)
-@pytest.mark.parametrize("database", [FIRESTORE_ENTERPRISE_DB], indirect=True)
-def test_pipeline_index_mode(database, query_docs):
-    """test pipeline query with explicit index mode"""
-
-    collection, _, allowed_vals = query_docs
-    client = collection._client
-    query = collection.where(filter=FieldFilter("a", "==", 1))
-    pipeline = client.pipeline().create_from(query)
-    with pytest.raises(InvalidArgument) as e:
-        pipeline.execute(index_mode="fake_index")
-    assert "Invalid index_mode: fake_index" in str(e)
-
-
 @pytest.mark.parametrize("database", TEST_DATABASES, indirect=True)
 def test_query_stream_w_read_time(query_docs, cleanup, database):
     collection, stored, allowed_vals = query_docs
@@ -1825,7 +1797,6 @@ def test_query_stream_w_read_time(query_docs, cleanup, database):
     assert new_values[new_ref.id] == new_data
 
 
-@pytest.mark.skipif(IS_KOKORO_TEST, reason="skipping pipeline verification on kokoro")
 @pytest.mark.parametrize("database", [FIRESTORE_ENTERPRISE_DB], indirect=True)
 def test_pipeline_w_read_time(query_docs, cleanup, database):
     collection, stored, allowed_vals = query_docs

tests/system/test_system_async.py

@@ -57,7 +57,6 @@
     ENTERPRISE_MODE_ERROR,
     TEST_DATABASES,
     TEST_DATABASES_W_ENTERPRISE,
-    IS_KOKORO_TEST,
     FIRESTORE_ENTERPRISE_DB,
 )
 
@@ -145,12 +144,6 @@ def _verify_explain_metrics_analyze_false(explain_metrics):
 
 @pytest.fixture(scope="session")
 def database(request):
-    from test__helpers import FIRESTORE_ENTERPRISE_DB
-
-    # enterprise mode currently does not support RunQuery calls in prod on kokoro test project
-    # TODO: remove skip when kokoro test project supports full enterprise mode
-    if request.param == FIRESTORE_ENTERPRISE_DB and IS_KOKORO_TEST:
-        pytest.skip("enterprise mode does not support RunQuery on kokoro")
     return request.param
 
 
@@ -181,11 +174,6 @@ async def verify_pipeline(query):
     """
     from google.cloud.firestore_v1.base_aggregation import BaseAggregationQuery
 
-    # return early on kokoro. Test project doesn't currently support pipelines
-    # TODO: enable pipeline verification when kokoro test project is whitelisted
-    if IS_KOKORO_TEST:
-        pytest.skip("skipping pipeline verification on kokoro")
-
     def _clean_results(results):
         if isinstance(results, dict):
             return {k: _clean_results(v) for k, v in results.items()}
@@ -1694,7 +1682,6 @@ async def test_pipeline_explain_options_using_additional_options(
     assert "Execution:" in text_stats
 
 
-@pytest.mark.skipif(IS_KOKORO_TEST, reason="skipping pipeline verification on kokoro")
 @pytest.mark.parametrize("database", [FIRESTORE_ENTERPRISE_DB], indirect=True)
 async def test_pipeline_w_read_time(query_docs, cleanup, database):
     collection, stored, allowed_vals = query_docs

tests/unit/v1/test_pipeline_result.py

@@ -213,23 +213,20 @@ def test_ctor(self):
         expected_transaction = object()
         expected_read_time = 123
         expected_explain_options = object()
-        expected_index_mode = "mode"
         expected_addtl_options = {}
         source = PipelineStream(
             expected_type,
             expected_pipeline,
             expected_transaction,
             expected_read_time,
             expected_explain_options,
-            expected_index_mode,
             expected_addtl_options,
         )
         instance = self._make_one(in_arr, source)
         assert instance._return_type == expected_type
         assert instance.pipeline == expected_pipeline
         assert instance._client == expected_pipeline._client
         assert instance._additonal_options == expected_addtl_options
-        assert instance._index_mode == expected_index_mode
         assert instance._explain_options == expected_explain_options
         assert instance._explain_stats is None
         assert instance._started is True
@@ -281,7 +278,6 @@ def _mock_init_args(self):
             "transaction": None,
             "read_time": None,
             "explain_options": None,
-            "index_mode": None,
             "additional_options": {},
         }
 
@@ -312,7 +308,6 @@ def test_explain_stats(self):
     @pytest.mark.parametrize(
         "init_kwargs,expected_options",
         [
-            ({"index_mode": "mode"}, {"index_mode": encode_value("mode")}),
             (
                 {"explain_options": PipelineExplainOptions()},
                 {"explain_options": encode_value({"mode": "analyze"})},
@@ -336,13 +331,6 @@ def test_explain_stats(self):
                 },
                 {"explain_options": encode_value("override")},
             ),
-            (
-                {
-                    "index_mode": "mode",
-                    "additional_options": {"index_mode": Constant("new")},
-                },
-                {"index_mode": encode_value("new")},
-            ),
         ],
     )
     def test_build_request_options(self, init_kwargs, expected_options):