From 7e055720096ad91ccf2733abbb47a00afa5ea6cb Mon Sep 17 00:00:00 2001
From: Afroz Alam <afroz.alam@snowflake.com>
Date: Wed, 17 Jul 2024 18:11:38 +0000
Subject: [PATCH] Fix telemetry test when sql simplifier is disabled (#1925)

---
 tests/integ/test_telemetry.py | 27 +++++++++++++++------------
 1 file changed, 15 insertions(+), 12 deletions(-)

diff --git a/tests/integ/test_telemetry.py b/tests/integ/test_telemetry.py
index 61debb4d131..e2f4f554eed 100644
--- a/tests/integ/test_telemetry.py
+++ b/tests/integ/test_telemetry.py
@@ -586,6 +586,8 @@ def test_execute_queries_api_calls(session, sql_simplifier_enabled):
     df.collect()
     # API calls don't change after query is executed
     query_plan_height = 2 if sql_simplifier_enabled else 3
+    filter = 1 if sql_simplifier_enabled else 2
+    low_impact = 3 if sql_simplifier_enabled else 2
 
     assert df._plan.api_calls == [
         {
@@ -594,8 +596,8 @@ def test_execute_queries_api_calls(session, sql_simplifier_enabled):
             "query_plan_height": query_plan_height,
             "query_plan_num_duplicate_nodes": 0,
             "query_plan_complexity": {
-                "filter": 1,
-                "low_impact": 3,
+                "filter": filter,
+                "low_impact": low_impact,
                 "function": 3,
                 "column": 3,
                 "literal": 5,
@@ -616,8 +618,8 @@ def test_execute_queries_api_calls(session, sql_simplifier_enabled):
             "query_plan_height": query_plan_height,
             "query_plan_num_duplicate_nodes": 0,
             "query_plan_complexity": {
-                "filter": 1,
-                "low_impact": 3,
+                "filter": filter,
+                "low_impact": low_impact,
                 "function": 3,
                 "column": 3,
                 "literal": 5,
@@ -638,8 +640,8 @@ def test_execute_queries_api_calls(session, sql_simplifier_enabled):
             "query_plan_height": query_plan_height,
             "query_plan_num_duplicate_nodes": 0,
             "query_plan_complexity": {
-                "filter": 1,
-                "low_impact": 3,
+                "filter": filter,
+                "low_impact": low_impact,
                 "function": 3,
                 "column": 3,
                 "literal": 5,
@@ -660,8 +662,8 @@ def test_execute_queries_api_calls(session, sql_simplifier_enabled):
             "query_plan_height": query_plan_height,
             "query_plan_num_duplicate_nodes": 0,
             "query_plan_complexity": {
-                "filter": 1,
-                "low_impact": 3,
+                "filter": filter,
+                "low_impact": low_impact,
                 "function": 3,
                 "column": 3,
                 "literal": 5,
@@ -682,8 +684,8 @@ def test_execute_queries_api_calls(session, sql_simplifier_enabled):
             "query_plan_height": query_plan_height,
             "query_plan_num_duplicate_nodes": 0,
             "query_plan_complexity": {
-                "filter": 1,
-                "low_impact": 3,
+                "filter": filter,
+                "low_impact": low_impact,
                 "function": 3,
                 "column": 3,
                 "literal": 5,
@@ -819,6 +821,7 @@ def test_dataframe_stat_functions_api_calls(session):
     # check to make sure that the original DF is unchanged
     assert df._plan.api_calls == [{"name": "Session.create_dataframe[values]"}]
 
+    column = 6 if session.sql_simplifier_enabled else 9
     crosstab = df.stat.crosstab("empid", "month")
     assert crosstab._plan.api_calls == [
         {
@@ -826,7 +829,7 @@ def test_dataframe_stat_functions_api_calls(session):
             "sql_simplifier_enabled": session.sql_simplifier_enabled,
             "query_plan_height": 4,
             "query_plan_num_duplicate_nodes": 0,
-            "query_plan_complexity": {"group_by": 1, "column": 6, "literal": 48},
+            "query_plan_complexity": {"group_by": 1, "column": column, "literal": 48},
         },
         {
             "name": "DataFrameStatFunctions.crosstab",
@@ -844,7 +847,7 @@ def test_dataframe_stat_functions_api_calls(session):
             "sql_simplifier_enabled": session.sql_simplifier_enabled,
             "query_plan_height": 4,
             "query_plan_num_duplicate_nodes": 0,
-            "query_plan_complexity": {"group_by": 1, "column": 6, "literal": 48},
+            "query_plan_complexity": {"group_by": 1, "column": column, "literal": 48},
         }
     ]