SNOW-2005919: update tests to fix disabled sql simplifier workflow (#3195)

Signed-off-by: Labanya Mukhopadhyay <[email protected]>

Author: sfc-gh-lmukhopadhyay

tests/integ/modin/groupby/test_groupby_transform.py

@@ -145,17 +145,7 @@ def test_dataframe_groupby_transform_conflicting_labels_negative():
         pd.show(df)
 
 
-@sql_count_checker(
-    query_count=7,
-    # Pivot step performs force materialization of the dataframe. So joins in subquery
-    # are not carried over but after removing the pivot step subquery is present
-    # multiple times resulting to increase in join count.
-    join_count=12,
-    udtf_count=1,
-    high_count_expected=True,
-    high_count_reason="performing two groupby transform operations that use UDTFs and compare with pandas",
-)
-def test_dataframe_groupby_transform_conflicting_labels():
+def test_dataframe_groupby_transform_conflicting_labels(session):
     """
     Based on SNOW-1361200 - The bug occurred because of conflicting UDTF columns appended during groupby transform
     operations in `create_udtf_for_groupby_apply`.
@@ -167,11 +157,19 @@ def transform_helper(df):
         df["A"] = df.groupby("X")["X_DATA"].transform("count")
         df["B"] = df.groupby("X")["X_DATA"].transform("count")
 
-    eval_snowpark_pandas_result(
-        *create_test_dfs({"X": [1, 2, 3, 1, 2, 2], "Y": [4, 5, 6, 7, 8, 9]}),
-        transform_helper,
-        inplace=True,
-    )
+    with SqlCounter(
+        query_count=7,
+        join_count=12 if session.sql_simplifier_enabled else 6,
+        udtf_count=1,
+        high_count_expected=True,
+        high_count_reason="performing two groupby transform operations that use UDTFs and compare "
+        "with pandas",
+    ):
+        eval_snowpark_pandas_result(
+            *create_test_dfs({"X": [1, 2, 3, 1, 2, 2], "Y": [4, 5, 6, 7, 8, 9]}),
+            transform_helper,
+            inplace=True,
+        )
 
 
 @sql_count_checker(

tests/integ/modin/test_df_to_snowpark_pandas.py

@@ -141,6 +141,10 @@ def test_to_snowpark_pandas_with_operations(
 def test_to_snowpark_pandas_duplicated_columns_raises(
     session, tmp_table_basic, relaxed_ordering
 ) -> None:
+    sql_simplifier_enabled_original = session.sql_simplifier_enabled
+    # Error is raised only when SQL simplifier is enabled.
+    session.sql_simplifier_enabled = True
+
     snowpark_df = session.table(tmp_table_basic)
     snowpark_df = snowpark_df.select(
         Column("ID"),
@@ -156,6 +160,7 @@ def test_to_snowpark_pandas_duplicated_columns_raises(
 
     with pytest.raises(SnowparkSQLException, match=pattern):
         snowpark_df.to_snowpark_pandas(relaxed_ordering=relaxed_ordering).head()
+    session.sql_simplifier_enabled = sql_simplifier_enabled_original
 
 
 @pytest.mark.parametrize("relaxed_ordering", [True, False])