NO SNOW: using mock instead of altering value in snowpark.context (#4025)

Author: sfc-gh-yuwang

tests/integ/test_df_aggregate.py

@@ -4,6 +4,7 @@
 #
 import decimal
 import math
+from unittest import mock
 
 import pytest
 
@@ -31,7 +32,6 @@
 )
 from snowflake.snowpark.mock._snowflake_data_type import ColumnEmulator, ColumnType
 from snowflake.snowpark.types import DoubleType, IntegerType, StructType, StructField
-import snowflake.snowpark.context as context
 from tests.utils import Utils
 
 
@@ -644,10 +644,9 @@ def test_agg_on_empty_df(session):
     reason="HAVING clause is not supported in local testing mode",
 )
 def test_agg_filter_snowpark_connect_compatible(session):
-    original_value = context._is_snowpark_connect_compatible_mode
-
-    try:
-        context._is_snowpark_connect_compatible_mode = True
+    with mock.patch(
+        "snowflake.snowpark.context._is_snowpark_connect_compatible_mode", True
+    ):
         df = session.create_dataframe(
             [(1, 2, 3), (3, 2, 1), (3, 2, 1)], ["a", "b", "c"]
         )
@@ -679,19 +678,16 @@ def test_agg_filter_snowpark_connect_compatible(session):
             df.filter(grouping("a") == 0).collect()
 
         Utils.check_answer(df.filter(col("a") > 1), [Row(3, 2, 1), Row(3, 2, 1)])
-    finally:
-        context._is_snowpark_connect_compatible_mode = original_value
 
 
 @pytest.mark.skipif(
     "config.getoption('local_testing_mode', default=False)",
     reason="ORDER BY append is not supported in local testing mode",
 )
 def test_agg_sort_snowpark_connect_compatible(session):
-    original_value = context._is_snowpark_connect_compatible_mode
-
-    try:
-        context._is_snowpark_connect_compatible_mode = True
+    with mock.patch(
+        "snowflake.snowpark.context._is_snowpark_connect_compatible_mode", True
+    ):
         df = session.create_dataframe(
             [(1, 2, 3), (3, 2, 1), (3, 2, 1)], ["a", "b", "c"]
         )
@@ -717,32 +713,27 @@ def test_agg_sort_snowpark_connect_compatible(session):
         # original behavior on dataframe without group by
         df4 = df.sort(col("a"))
         Utils.check_answer(df4, [Row(1, 2, 3), Row(3, 2, 1), Row(3, 2, 1)])
-    finally:
-        context._is_snowpark_connect_compatible_mode = original_value
 
 
 def test_agg_no_grouping_exprs_limit_snowpark_connect_compatible(session):
-    original_value = context._is_snowpark_connect_compatible_mode
-    try:
-        context._is_snowpark_connect_compatible_mode = True
+    with mock.patch(
+        "snowflake.snowpark.context._is_snowpark_connect_compatible_mode", True
+    ):
         df = session.create_dataframe([[1, 2], [3, 4], [1, 4]], schema=["A", "B"])
         result = df.agg(sum_(col("a"))).limit(2)
         Utils.check_answer(result, [Row(5)])
         result = df.group_by().agg(sum_(col("b"))).limit(2)
         Utils.check_answer(result, [Row(10)])
-    finally:
-        context._is_snowpark_connect_compatible_mode = original_value
 
 
 @pytest.mark.skipif(
     "config.getoption('local_testing_mode', default=False)",
     reason="HAVING and ORDER BY append are not supported in local testing mode",
 )
 def test_agg_filter_and_sort_with_grouping_snowpark_connect_compatible(session):
-    original_value = context._is_snowpark_connect_compatible_mode
-
-    try:
-        context._is_snowpark_connect_compatible_mode = True
+    with mock.patch(
+        "snowflake.snowpark.context._is_snowpark_connect_compatible_mode", True
+    ):
         df = session.create_dataframe(
             [
                 ("dotNET", 2012, 10000),
@@ -852,19 +843,16 @@ def test_agg_filter_and_sort_with_grouping_snowpark_connect_compatible(session):
         # First row should have highest grouping value (1)
         results6 = df6.collect()
         assert results6[0][2] == 1  # gc=1 for NULL course
-    finally:
-        context._is_snowpark_connect_compatible_mode = original_value
 
 
 @pytest.mark.skipif(
     "config.getoption('local_testing_mode', default=False)",
     reason="HAVING, ORDER BY append, and limit append are not supported in local testing mode",
 )
 def test_filter_sort_limit_snowpark_connect_compatible(session, sql_simplifier_enabled):
-    original_value = context._is_snowpark_connect_compatible_mode
-
-    try:
-        context._is_snowpark_connect_compatible_mode = True
+    with mock.patch(
+        "snowflake.snowpark.context._is_snowpark_connect_compatible_mode", True
+    ):
         df = session.create_dataframe(
             [(1, 2, 3), (3, 2, 1), (3, 2, 1)], ["a", "b", "c"]
         )
@@ -939,9 +927,6 @@ def test_filter_sort_limit_snowpark_connect_compatible(session, sql_simplifier_e
         # Should have 4 SELECT statements
         assert query6.upper().count("SELECT") == 4 if sql_simplifier_enabled else 5
 
-    finally:
-        context._is_snowpark_connect_compatible_mode = original_value
-
 
 @pytest.mark.skipif(
     "config.getoption('local_testing_mode', default=False)",

tests/integ/test_udtf.py

@@ -8,10 +8,11 @@
 import sys
 from textwrap import dedent
 from typing import Dict, List, Tuple
+from unittest import mock
 
 import pytest
 
-from snowflake.snowpark import Row, Table, context
+from snowflake.snowpark import Row, Table
 from snowflake.snowpark._internal.analyzer.analyzer_utils import unquote_if_quoted
 from snowflake.snowpark._internal.utils import TempObjectType
 from snowflake.snowpark.exceptions import SnowparkSQLException
@@ -557,100 +558,102 @@ def __init__(self) -> None:
         df._column_map.columns = [Column("id"), Column("v")]
         return df
 
-    context._is_snowpark_connect_compatible_mode = True
-
-    df = create_snowpark_compatible_dataframe()
+    with mock.patch(
+        "snowflake.snowpark.context._is_snowpark_connect_compatible_mode", True
+    ):
 
-    def normalize(pdf):
-        v = pdf.v
-        return pdf.assign(v=(v - v.mean()) / v.std())
-
-    df = (
-        df.group_by("id")
-        .applyInPandas(
-            normalize,
-            output_schema=StructType(
-                [
-                    StructField("id", IntegerType()),
-                    StructField("v", DoubleType()),
-                ]
-            ),
+        df = create_snowpark_compatible_dataframe()
+
+        def normalize(pdf):
+            v = pdf.v
+            return pdf.assign(v=(v - v.mean()) / v.std())
+
+        df = (
+            df.group_by("id")
+            .applyInPandas(
+                normalize,
+                output_schema=StructType(
+                    [
+                        StructField("id", IntegerType()),
+                        StructField("v", DoubleType()),
+                    ]
+                ),
+            )
+            .orderBy(["id", "v"])
         )
-        .orderBy(["id", "v"])
-    )
-
-    Utils.check_answer(
-        df,
-        [
-            Row(ID=1, V=-0.7071067811865475),
-            Row(ID=1, V=0.7071067811865475),
-            Row(ID=2, V=-0.8320502943378437),
-            Row(ID=2, V=-0.2773500981126146),
-            Row(ID=2, V=1.1094003924504583),
-        ],
-    )
 
-    df = create_snowpark_compatible_dataframe()
-
-    def sum_func(key, pdf):
-        # key is a tuple of two numpy.int64s, which is the values
-        # of 'id' and 'ceil(df.v / 2)' for the current group
-        return pd.DataFrame([key + (pdf.v.sum(),)])
-
-    df = (
-        df.group_by("id", ceil(df.v / 2).alias("newcol"))
-        .applyInPandas(
-            sum_func,
-            output_schema=StructType(
-                [
-                    StructField("id", IntegerType()),
-                    StructField("c", IntegerType()),
-                    StructField("v", DoubleType()),
-                ]
-            ),
+        Utils.check_answer(
+            df,
+            [
+                Row(ID=1, V=-0.7071067811865475),
+                Row(ID=1, V=0.7071067811865475),
+                Row(ID=2, V=-0.8320502943378437),
+                Row(ID=2, V=-0.2773500981126146),
+                Row(ID=2, V=1.1094003924504583),
+            ],
         )
-        .orderBy(["id", "v"])
-    )
-
-    Utils.check_answer(
-        df,
-        [
-            Row(ID=1, C=1, V=3.0),
-            Row(ID=2, C=2, V=3.0),
-            Row(ID=2, C=3, V=5.0),
-            Row(ID=2, C=5, V=10.0),
-        ],
-    )
 
-    df = create_snowpark_compatible_dataframe()
+        df = create_snowpark_compatible_dataframe()
+
+        def sum_func(key, pdf):
+            # key is a tuple of two numpy.int64s, which is the values
+            # of 'id' and 'ceil(df.v / 2)' for the current group
+            return pd.DataFrame([key + (pdf.v.sum(),)])
+
+        df = (
+            df.group_by("id", ceil(df.v / 2).alias("newcol"))
+            .applyInPandas(
+                sum_func,
+                output_schema=StructType(
+                    [
+                        StructField("id", IntegerType()),
+                        StructField("c", IntegerType()),
+                        StructField("v", DoubleType()),
+                    ]
+                ),
+            )
+            .orderBy(["id", "v"])
+        )
 
-    def sum_func_with_single_input(pdf):
-        # key is a tuple of two numpy.int64s, which is the values
-        # of 'id' and 'ceil(df.v / 2)' for the current group
-        return pd.DataFrame([(pdf.v.sum(),)])
+        Utils.check_answer(
+            df,
+            [
+                Row(ID=1, C=1, V=3.0),
+                Row(ID=2, C=2, V=3.0),
+                Row(ID=2, C=3, V=5.0),
+                Row(ID=2, C=5, V=10.0),
+            ],
+        )
 
-    df = (
-        df.group_by("id", ceil(df.v / 2))
-        .applyInPandas(
-            sum_func_with_single_input,
-            output_schema=StructType(
-                [
-                    StructField("v", DoubleType()),
-                ]
-            ),
+        df = create_snowpark_compatible_dataframe()
+
+        def sum_func_with_single_input(pdf):
+            # key is a tuple of two numpy.int64s, which is the values
+            # of 'id' and 'ceil(df.v / 2)' for the current group
+            return pd.DataFrame([(pdf.v.sum(),)])
+
+        df = (
+            df.group_by("id", ceil(df.v / 2))
+            .applyInPandas(
+                sum_func_with_single_input,
+                output_schema=StructType(
+                    [
+                        StructField("v", DoubleType()),
+                    ]
+                ),
+            )
+            .orderBy(["v"])
         )
-        .orderBy(["v"])
-    )
 
-    Utils.check_answer(
-        df,
-        [
-            Row(V=3.0),
-            Row(V=3.0),
-            Row(V=5.0),
-            Row(V=10.0),
-        ],
-    )
+        Utils.check_answer(
+            df,
+            [
+                Row(V=3.0),
+                Row(V=3.0),
+                Row(V=5.0),
+                Row(V=10.0),
+            ],
+        )
 
 
 @pytest.mark.skipif(IS_IN_STORED_PROC, reason="Cannot create session in SP")