From e7dda7c9a75a9c14a5176d2e359181965d73885f Mon Sep 17 00:00:00 2001
From: Linchin <lingqing.gan@gmail.com>
Date: Thu, 25 Jul 2024 23:49:28 -0700
Subject: [PATCH] add system tests

---
 tests/system/test_system.py | 94 ++++++++++++++++++++++++++++++++-----
 1 file changed, 81 insertions(+), 13 deletions(-)

diff --git a/tests/system/test_system.py b/tests/system/test_system.py
index 3b30d2e38..454a5ce00 100644
--- a/tests/system/test_system.py
+++ b/tests/system/test_system.py
@@ -862,17 +862,37 @@ def test_query_stream_w_offset(query_docs, database):
     FIRESTORE_EMULATOR, reason="Query profile not supported in emulator."
 )
 @pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True)
-def test_query_stream_w_explain_options(query_docs, database):
+def test_query_stream_w_no_explain_options(query_docs, database):
+    from google.cloud.firestore_v1.query_profile import QueryExplainError
+
+    collection, _, allowed_vals = query_docs
+    num_vals = len(allowed_vals)
+    query = collection.where(filter=FieldFilter("a", "in", [1, num_vals + 100]))
+
+    # If no explain_option is passed, raise an exception if explain_metrics
+    # is called
+    results = query.stream()
+    with pytest.raises(QueryExplainError, match="explain_options not set on query"):
+        results.explain_metrics
+
+
+@pytest.mark.skipif(
+    FIRESTORE_EMULATOR, reason="Query profile not supported in emulator."
+)
+@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True)
+def test_query_stream_w_explain_options_analyze_true(query_docs, database):
     from google.cloud.firestore_v1.query_profile import (
+        ExecutionStats,
         ExplainMetrics,
         ExplainOptions,
+        PlanSummary,
         QueryExplainError,
     )
 
-    collection, stored, allowed_vals = query_docs
+    collection, _, allowed_vals = query_docs
     num_vals = len(allowed_vals)
     query = collection.where(filter=FieldFilter("a", "in", [1, num_vals + 100]))
-    results_1 = query.stream(explain_options=ExplainOptions(analyze=True))
+    results = query.stream(explain_options=ExplainOptions(analyze=True))
 
     # An exception should be raised when accessing explain_metrics before query
     # finishes.
@@ -880,17 +900,66 @@ def test_query_stream_w_explain_options(query_docs, database):
         QueryExplainError,
         match="explain_metrics not available until query is complete",
     ):
-        results_1.explain_metrics
+        results.explain_metrics
+
+    # Finish iterating results, and explain_metrics should be available.
+    num_results = len(list(results))
+    assert isinstance(results.explain_metrics, ExplainMetrics)
+
+    # verify plan_summary
+    plan_summary = results.explain_metrics.plan_summary
+    assert isinstance(plan_summary, PlanSummary)
+    assert len(plan_summary.indexes_used) > 0
+    assert plan_summary.indexes_used[0]["properties"] == "(a ASC, __name__ ASC)"
+    assert plan_summary.indexes_used[0]["query_scope"] == "Collection"
+
+    # verify execution_stats
+    execution_stats = results.explain_metrics.execution_stats
+    assert isinstance(execution_stats, ExecutionStats)
+    assert execution_stats.results_returned == num_results
+    assert execution_stats.read_operations == num_results
+    duration = execution_stats.execution_duration.total_seconds()
+    assert duration > 0
+    assert duration < 1  # we expect a number closer to 0.05
+    assert isinstance(execution_stats.debug_stats, dict)
+    assert "billing_details" in execution_stats.debug_stats
+    assert "documents_scanned" in execution_stats.debug_stats
+    assert "index_entries_scanned" in execution_stats.debug_stats
+    assert len(execution_stats.debug_stats) > 0
 
-    # Finish iterating results, and explain_metrics should be available
-    list(results_1)
-    assert isinstance(results_1.explain_metrics, ExplainMetrics)
 
-    # If no explain_option is passed, raise an exception if explain_metrics
-    # is called
-    results_2 = query.stream()
-    with pytest.raises(QueryExplainError, match="explain_options not set on query"):
-        results_2.explain_metrics
+@pytest.mark.skipif(
+    FIRESTORE_EMULATOR, reason="Query profile not supported in emulator."
+)
+@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True)
+def test_query_stream_w_explain_options_analyze_false(query_docs, database):
+    from google.cloud.firestore_v1.query_profile import (
+        ExplainMetrics,
+        ExplainOptions,
+        PlanSummary,
+        QueryExplainError,
+    )
+
+    collection, _, allowed_vals = query_docs
+    num_vals = len(allowed_vals)
+    query = collection.where(filter=FieldFilter("a", "in", [1, num_vals + 100]))
+    results = query.stream(explain_options=ExplainOptions(analyze=False))
+
+    assert isinstance(results.explain_metrics, ExplainMetrics)
+
+    # verify plan_summary
+    plan_summary = results.explain_metrics.plan_summary
+    assert isinstance(plan_summary, PlanSummary)
+    assert len(plan_summary.indexes_used) > 0
+    assert plan_summary.indexes_used[0]["properties"] == "(a ASC, __name__ ASC)"
+    assert plan_summary.indexes_used[0]["query_scope"] == "Collection"
+
+    # verify execution_stats isn't available
+    with pytest.raises(
+        QueryExplainError,
+        match="execution_stats not available when explain_options.analyze=False",
+    ):
+        results.explain_metrics.execution_stats
 
 
 @pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True)
@@ -2424,7 +2493,6 @@ def test_or_query_in_transaction(client, cleanup, database):
         def in_transaction(transaction):
             global inner_fn_ran
             result = query.get(transaction=transaction)
-
             assert len(result) == 2
             # both documents should have a == 1
             assert result[0].get("a") == 1