use subtests for verify_pipeline

Author: daniel-sanche

packages/google-cloud-firestore/noxfile.py

@@ -71,7 +71,7 @@
 SYSTEM_TEST_PYTHON_VERSIONS: List[str] = ALL_PYTHON
 SYSTEM_TEST_STANDARD_DEPENDENCIES = [
     "mock",
-    "pytest",
+    "pytest>9.0",
     "google-cloud-testutils",
 ]
 SYSTEM_TEST_EXTERNAL_DEPENDENCIES: List[str] = [

packages/google-cloud-firestore/tests/system/test_system.py

@@ -83,69 +83,73 @@ def cleanup():
     for operation in operations:
         operation()
 
-
-def verify_pipeline(query):
+@pytest.fixture
+def verify_pipeline(subtests):
     """
-    This function ensures a pipeline produces the same
-    results as the query it is derived from
+    This fixture provide a subtest function which
+    ensures a pipeline produces the same results as the query it is derived
+    from
 
     It can be attached to existing query tests to check both
     modalities at the same time
 
     Pipelines are only supported on enterprise dbs. Skip other environments
     """
-    from google.cloud.firestore_v1.base_aggregation import BaseAggregationQuery
-
-    client = query._client
-    if FIRESTORE_EMULATOR:
-        print("skip pipeline verification on emulator")
-        return
-    if client._database != FIRESTORE_ENTERPRISE_DB:
-        print("pipelines only supports enterprise db")
-        return
-
-    def _clean_results(results):
-        if isinstance(results, dict):
-            return {k: _clean_results(v) for k, v in results.items()}
-        elif isinstance(results, list):
-            return [_clean_results(r) for r in results]
-        elif isinstance(results, float) and math.isnan(results):
-            return "__NAN_VALUE__"
-        else:
-            return results
 
-    query_exception = None
-    query_results = None
-    try:
-        try:
-            if isinstance(query, BaseAggregationQuery):
-                # aggregation queries return a list of lists of aggregation results
-                query_results = _clean_results(
-                    list(
-                        itertools.chain.from_iterable(
-                            [[a._to_dict() for a in s] for s in query.get()]
+    def _verifier(query):
+        from google.cloud.firestore_v1.base_aggregation import BaseAggregationQuery
+        with subtests.test(msg="verify_pipeline"):
+
+            client = query._client
+            if FIRESTORE_EMULATOR:
+                pytest.skip("skip pipeline verification on emulator")
+            if client._database != FIRESTORE_ENTERPRISE_DB:
+                pytest.skip("pipelines only supports enterprise db")
+
+            def _clean_results(results):
+                if isinstance(results, dict):
+                    return {k: _clean_results(v) for k, v in results.items()}
+                elif isinstance(results, list):
+                    return [_clean_results(r) for r in results]
+                elif isinstance(results, float) and math.isnan(results):
+                    return "__NAN_VALUE__"
+                else:
+                    return results
+
+            query_exception = None
+            query_results = None
+            try:
+                try:
+                    if isinstance(query, BaseAggregationQuery):
+                        # aggregation queries return a list of lists of aggregation results
+                        query_results = _clean_results(
+                            list(
+                                itertools.chain.from_iterable(
+                                    [[a._to_dict() for a in s] for s in query.get()]
+                                )
+                            )
                         )
-                    )
-                )
-            else:
-                # other qureies return a simple list of results
-                query_results = _clean_results([s.to_dict() for s in query.get()])
-        except Exception as e:
-            # if we expect the query to fail, capture the exception
-            query_exception = e
-        pipeline = client.pipeline().create_from(query)
-        if query_exception:
-            # ensure that the pipeline uses same error as query
-            with pytest.raises(query_exception.__class__):
-                pipeline.execute()
-        else:
-            # ensure results match query
-            pipeline_results = _clean_results([s.data() for s in pipeline.execute()])
-            assert query_results == pipeline_results
-    except FailedPrecondition as e:
-        # if testing against a non-enterprise db, skip this check
-        if ENTERPRISE_MODE_ERROR not in e.message:
-            raise e
+                    else:
+                        # other qureies return a simple list of results
+                        query_results = _clean_results([s.to_dict() for s in query.get()])
+                except Exception as e:
+                    # if we expect the query to fail, capture the exception
+                    query_exception = e
+                pipeline = client.pipeline().create_from(query)
+                if query_exception:
+                    # ensure that the pipeline uses same error as query
+                    with pytest.raises(query_exception.__class__):
+                        pipeline.execute()
+                else:
+                    # ensure results match query
+                    pipeline_results = _clean_results([s.data() for s in pipeline.execute()])
+                    assert query_results == pipeline_results
+            except FailedPrecondition as e:
+                # if testing against a non-enterprise db, skip this check
+                if ENTERPRISE_MODE_ERROR not in e.message:
+                    raise e
+
+    return _verifier
 
 
 @pytest.mark.parametrize("database", TEST_DATABASES, indirect=True)
@@ -1300,7 +1304,7 @@ def query(collection):
 
 
 @pytest.mark.parametrize("database", TEST_DATABASES_W_ENTERPRISE, indirect=True)
-def test_query_stream_legacy_where(query_docs, database):
+def test_query_stream_legacy_where(query_docs, database, verify_pipeline):
     """Assert the legacy code still works and returns value"""
     collection, stored, allowed_vals = query_docs
     with pytest.warns(
@@ -1317,7 +1321,7 @@ def test_query_stream_legacy_where(query_docs, database):
 
 
 @pytest.mark.parametrize("database", TEST_DATABASES_W_ENTERPRISE, indirect=True)
-def test_query_stream_w_simple_field_eq_op(query_docs, database):
+def test_query_stream_w_simple_field_eq_op(query_docs, database, verify_pipeline):
     collection, stored, allowed_vals = query_docs
     query = collection.where(filter=FieldFilter("a", "==", 1))
     values = {snapshot.id: snapshot.to_dict() for snapshot in query.stream()}
@@ -1329,7 +1333,7 @@ def test_query_stream_w_simple_field_eq_op(query_docs, database):
 
 
 @pytest.mark.parametrize("database", TEST_DATABASES_W_ENTERPRISE, indirect=True)
-def test_query_stream_w_simple_field_array_contains_op(query_docs, database):
+def test_query_stream_w_simple_field_array_contains_op(query_docs, database, verify_pipeline):
     collection, stored, allowed_vals = query_docs
     query = collection.where(filter=FieldFilter("c", "array_contains", 1))
     values = {snapshot.id: snapshot.to_dict() for snapshot in query.stream()}
@@ -1341,7 +1345,7 @@ def test_query_stream_w_simple_field_array_contains_op(query_docs, database):
 
 
 @pytest.mark.parametrize("database", TEST_DATABASES_W_ENTERPRISE, indirect=True)
-def test_query_stream_w_simple_field_in_op(query_docs, database):
+def test_query_stream_w_simple_field_in_op(query_docs, database, verify_pipeline):
     collection, stored, allowed_vals = query_docs
     num_vals = len(allowed_vals)
     query = collection.where(filter=FieldFilter("a", "in", [1, num_vals + 100]))
@@ -1354,7 +1358,7 @@ def test_query_stream_w_simple_field_in_op(query_docs, database):
 
 
 @pytest.mark.parametrize("database", TEST_DATABASES_W_ENTERPRISE, indirect=True)
-def test_query_stream_w_not_eq_op(query_docs, database):
+def test_query_stream_w_not_eq_op(query_docs, database, verify_pipeline):
     collection, stored, allowed_vals = query_docs
     query = collection.where(filter=FieldFilter("stats.sum", "!=", 4))
     values = {snapshot.id: snapshot.to_dict() for snapshot in query.stream()}
@@ -1377,7 +1381,7 @@ def test_query_stream_w_not_eq_op(query_docs, database):
 
 
 @pytest.mark.parametrize("database", TEST_DATABASES_W_ENTERPRISE, indirect=True)
-def test_query_stream_w_simple_not_in_op(query_docs, database):
+def test_query_stream_w_simple_not_in_op(query_docs, database, verify_pipeline):
     collection, stored, allowed_vals = query_docs
     num_vals = len(allowed_vals)
     query = collection.where(
@@ -1390,7 +1394,7 @@ def test_query_stream_w_simple_not_in_op(query_docs, database):
 
 
 @pytest.mark.parametrize("database", TEST_DATABASES_W_ENTERPRISE, indirect=True)
-def test_query_stream_w_simple_field_array_contains_any_op(query_docs, database):
+def test_query_stream_w_simple_field_array_contains_any_op(query_docs, database, verify_pipeline):
     collection, stored, allowed_vals = query_docs
     num_vals = len(allowed_vals)
     query = collection.where(
@@ -1405,7 +1409,7 @@ def test_query_stream_w_simple_field_array_contains_any_op(query_docs, database)
 
 
 @pytest.mark.parametrize("database", TEST_DATABASES_W_ENTERPRISE, indirect=True)
-def test_query_stream_w_order_by(query_docs, database):
+def test_query_stream_w_order_by(query_docs, database, verify_pipeline):
     collection, stored, allowed_vals = query_docs
     query = collection.order_by("b", direction=firestore.Query.DESCENDING)
     values = [(snapshot.id, snapshot.to_dict()) for snapshot in query.stream()]
@@ -1420,7 +1424,7 @@ def test_query_stream_w_order_by(query_docs, database):
 
 
 @pytest.mark.parametrize("database", TEST_DATABASES_W_ENTERPRISE, indirect=True)
-def test_query_stream_w_field_path(query_docs, database):
+def test_query_stream_w_field_path(query_docs, database, verify_pipeline):
     collection, stored, allowed_vals = query_docs
     query = collection.where(filter=FieldFilter("stats.sum", ">", 4))
     values = {snapshot.id: snapshot.to_dict() for snapshot in query.stream()}
@@ -1459,7 +1463,7 @@ def test_query_stream_w_start_end_cursor(query_docs, database):
 
 
 @pytest.mark.parametrize("database", TEST_DATABASES_W_ENTERPRISE, indirect=True)
-def test_query_stream_wo_results(query_docs, database):
+def test_query_stream_wo_results(query_docs, database, verify_pipeline):
     collection, stored, allowed_vals = query_docs
     num_vals = len(allowed_vals)
     query = collection.where(filter=FieldFilter("b", "==", num_vals + 100))
@@ -1486,7 +1490,7 @@ def test_query_stream_w_projection(query_docs, database):
 
 
 @pytest.mark.parametrize("database", TEST_DATABASES_W_ENTERPRISE, indirect=True)
-def test_query_stream_w_multiple_filters(query_docs, database):
+def test_query_stream_w_multiple_filters(query_docs, database, verify_pipeline):
     collection, stored, allowed_vals = query_docs
     query = collection.where(filter=FieldFilter("stats.product", ">", 5)).where(
         filter=FieldFilter("stats.product", "<", 10)
@@ -1507,7 +1511,7 @@ def test_query_stream_w_multiple_filters(query_docs, database):
 
 
 @pytest.mark.parametrize("database", TEST_DATABASES_W_ENTERPRISE, indirect=True)
-def test_query_stream_w_offset(query_docs, database):
+def test_query_stream_w_offset(query_docs, database, verify_pipeline):
     collection, stored, allowed_vals = query_docs
     num_vals = len(allowed_vals)
     offset = 3
@@ -1528,7 +1532,7 @@ def test_query_stream_w_offset(query_docs, database):
 )
 @pytest.mark.parametrize("method", ["stream", "get"])
 @pytest.mark.parametrize("database", TEST_DATABASES_W_ENTERPRISE, indirect=True)
-def test_query_stream_or_get_w_no_explain_options(query_docs, database, method):
+def test_query_stream_or_get_w_no_explain_options(query_docs, database, method, verify_pipeline):
     from google.cloud.firestore_v1.query_profile import QueryExplainError
 
     collection, _, allowed_vals = query_docs
@@ -1892,7 +1896,7 @@ def test_query_with_order_dot_key(client, cleanup, database):
 
 
 @pytest.mark.parametrize("database", TEST_DATABASES, indirect=True)
-def test_query_unary(client, cleanup, database):
+def test_query_unary(client, cleanup, database, verify_pipeline):
     collection_name = "unary" + UNIQUE_RESOURCE_ID
     collection = client.collection(collection_name)
     field_name = "foo"
@@ -1949,7 +1953,7 @@ def test_query_unary(client, cleanup, database):
 
 
 @pytest.mark.parametrize("database", TEST_DATABASES_W_ENTERPRISE, indirect=True)
-def test_collection_group_queries(client, cleanup, database):
+def test_collection_group_queries(client, cleanup, database, verify_pipeline):
     collection_group = "b" + UNIQUE_RESOURCE_ID
 
     doc_paths = [
@@ -2026,7 +2030,7 @@ def test_collection_group_queries_startat_endat(client, cleanup, database):
 
 
 @pytest.mark.parametrize("database", TEST_DATABASES_W_ENTERPRISE, indirect=True)
-def test_collection_group_queries_filters(client, cleanup, database):
+def test_collection_group_queries_filters(client, cleanup, database, verify_pipeline):
     collection_group = "b" + UNIQUE_RESOURCE_ID
 
     doc_paths = [
@@ -2817,7 +2821,7 @@ def on_snapshot(docs, changes, read_time):
 
 
 @pytest.mark.parametrize("database", TEST_DATABASES_W_ENTERPRISE, indirect=True)
-def test_repro_429(client, cleanup, database):
+def test_repro_429(client, cleanup, database, verify_pipeline):
     # See: https://github.com/googleapis/python-firestore/issues/429
     now = datetime.datetime.now(tz=datetime.timezone.utc)
     collection = client.collection("repro-429" + UNIQUE_RESOURCE_ID)
@@ -3412,7 +3416,7 @@ def test_aggregation_query_stream_or_get_w_explain_options_analyze_false(
 
 
 @pytest.mark.parametrize("database", TEST_DATABASES_W_ENTERPRISE, indirect=True)
-def test_query_with_and_composite_filter(collection, database):
+def test_query_with_and_composite_filter(collection, database, verify_pipeline):
     and_filter = And(
         filters=[
             FieldFilter("stats.product", ">", 5),
@@ -3428,7 +3432,7 @@ def test_query_with_and_composite_filter(collection, database):
 
 
 @pytest.mark.parametrize("database", TEST_DATABASES_W_ENTERPRISE, indirect=True)
-def test_query_with_or_composite_filter(collection, database):
+def test_query_with_or_composite_filter(collection, database, verify_pipeline):
     or_filter = Or(
         filters=[
             FieldFilter("stats.product", ">", 5),
@@ -3462,6 +3466,7 @@ def test_aggregation_queries_with_read_time(
     database,
     aggregation_type,
     expected_value,
+    verify_pipeline,
 ):
     """
     Ensure that all aggregation queries work when read_time is passed into
@@ -3500,7 +3505,7 @@ def test_aggregation_queries_with_read_time(
 
 
 @pytest.mark.parametrize("database", TEST_DATABASES_W_ENTERPRISE, indirect=True)
-def test_query_with_complex_composite_filter(collection, database):
+def test_query_with_complex_composite_filter(collection, database, verify_pipeline):
     field_filter = FieldFilter("b", "==", 0)
     or_filter = Or(
         filters=[FieldFilter("stats.sum", "==", 0), FieldFilter("stats.sum", "==", 4)]
@@ -3558,6 +3563,7 @@ def test_aggregation_query_in_transaction(
     aggregation_type,
     aggregation_args,
     expected,
+    verify_pipeline,
 ):
     """
     Test creating an aggregation query inside a transaction
@@ -3599,7 +3605,7 @@ def in_transaction(transaction):
 
 
 @pytest.mark.parametrize("database", TEST_DATABASES_W_ENTERPRISE, indirect=True)
-def test_or_query_in_transaction(client, cleanup, database):
+def test_or_query_in_transaction(client, cleanup, database, verify_pipeline):
     """
     Test running or query inside a transaction. Should pass transaction id along with request
     """