Fix query counts

snowflakedb · Jun 25, 2024 · f16212c · f16212c
1 parent 501875b
commit f16212c
Showing 1 changed file with 12 additions and 12 deletions.
diff --git a/tests/integ/modin/frame/test_assign.py b/tests/integ/modin/frame/test_assign.py
@@ -17,7 +17,7 @@
 )
 
 
-@sql_count_checker(query_count=7, join_count=1)
+@sql_count_checker(query_count=2, join_count=1)
 def test_assign_basic_series():
     snow_df, native_df = create_test_dfs(
         [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
@@ -36,7 +36,7 @@ def assign_func(df):
     eval_snowpark_pandas_result(snow_df, native_df, assign_func)
 
 
-@sql_count_checker(query_count=7, join_count=1)
+@sql_count_checker(query_count=2, join_count=1)
 @pytest.mark.parametrize(
     "index", [[2, 1, 0], [4, 5, 6]], ids=["reversed_index", "different_index"]
 )
@@ -61,7 +61,7 @@ def assign_func(df):
 @pytest.mark.parametrize("new_col_value", [2, [10, 11, 12], "x"])
 def test_assign_basic_non_pandas_object(new_col_value):
     join_count = 2 if isinstance(new_col_value, list) else 0
-    with SqlCounter(query_count=7, join_count=join_count):
+    with SqlCounter(query_count=2, join_count=join_count):
         snow_df, native_df = create_test_dfs(
             [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
             columns=native_pd.Index(list("abc"), name="columns"),
@@ -74,7 +74,7 @@ def test_assign_basic_non_pandas_object(new_col_value):
         )
 
 
-@sql_count_checker(query_count=7, join_count=2)
+@sql_count_checker(query_count=2, join_count=2)
 def test_assign_invalid_long_column_length_negative():
     # pandas errors out in this test, since we are attempting to assign a column of length 5 to a DataFrame with length 3.
     # Snowpark pandas on the other hand, just truncates the last element of the new column so that it is the correct length. If we wanted
@@ -98,7 +98,7 @@ def test_assign_invalid_long_column_length_negative():
     assert_snowpark_pandas_equals_to_pandas_without_dtypecheck(snow_df, native_df)
 
 
-@sql_count_checker(query_count=7, join_count=2)
+@sql_count_checker(query_count=2, join_count=2)
 def test_assign_invalid_short_column_length_negative():
     # pandas errors out in this test, since we are attempting to assign a column of length 2 to a DataFrame with length 3.
     # Snowpark pandas on the other hand, just broadcasts the last element of the new column so that it is filled. If we wanted
@@ -122,7 +122,7 @@ def test_assign_invalid_short_column_length_negative():
     assert_snowpark_pandas_equals_to_pandas_without_dtypecheck(snow_df, native_df)
 
 
-@sql_count_checker(query_count=7, join_count=1)
+@sql_count_checker(query_count=2, join_count=1)
 def test_assign_short_series():
     snow_df, native_df = create_test_dfs(
         [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
@@ -136,7 +136,7 @@ def test_assign_short_series():
     assert_snowpark_pandas_equals_to_pandas_without_dtypecheck(snow_df, native_df)
 
 
-@sql_count_checker(query_count=7, join_count=1)
+@sql_count_checker(query_count=2, join_count=1)
 @pytest.mark.parametrize(
     "index", [[1, 0], [4, 5]], ids=["reversed_index", "different_index"]
 )
@@ -153,7 +153,7 @@ def test_assign_short_series_mismatched_index(index):
     assert_snowpark_pandas_equals_to_pandas_without_dtypecheck(snow_df, native_df)
 
 
-@sql_count_checker(query_count=7)
+@sql_count_checker(query_count=2)
 @pytest.mark.parametrize(
     "callable_fn",
     [lambda x: x["a"], lambda x: x["a"] + x["b"]],
@@ -172,7 +172,7 @@ def test_assign_basic_callable(callable_fn):
     )
 
 
-@sql_count_checker(query_count=7)
+@sql_count_checker(query_count=2)
 def test_assign_chained_callable():
     snow_df, native_df = create_test_dfs(
         [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
@@ -190,7 +190,7 @@ def test_assign_chained_callable():
     )
 
 
-@sql_count_checker(query_count=6)
+@sql_count_checker(query_count=1)
 def test_assign_chained_callable_wrong_order():
     snow_df, native_df = create_test_dfs(
         [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
@@ -212,7 +212,7 @@ def test_assign_chained_callable_wrong_order():
     )
 
 
-@sql_count_checker(query_count=7)
+@sql_count_checker(query_count=2)
 def test_assign_self_columns():
     snow_df, native_df = create_test_dfs(
         [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
@@ -226,7 +226,7 @@ def test_assign_self_columns():
     )
 
 
-@sql_count_checker(query_count=7, join_count=2)
+@sql_count_checker(query_count=2, join_count=2)
 def test_overwrite_columns_via_assign():
     snow_df, native_df = create_test_dfs(
         [[1, 2, 3], [4, 5, 6], [7, 8, 9]],