Local Testing Framework¶

Prerequisites¶

To use the local testing framework:

• You must use version 1.11.1 or higher of the Snowpark Python library with the optional dependency pandas. Install by running pip install "snowflake-snowpark-python[pandas]"

• The supported versions of Python are:

  • 3.8

  • 3.9

  • 3.10

  • 3.11

Creating a Session and Enabling Local Testing¶

To get started, create a Snowpark Session and set the local testing configuration to True.

from snowflake.snowpark import Session

session = Session.builder.config('local_testing', True).create()

After the session is created, you can use it to create and operate on DataFrames.

df = session.create_dataframe([[1,2],[3,4]],['a','b'])
df.with_column('c', df['a']+df['b']).show()

Loading Data¶

You can create Snowpark DataFrames from Python primitives, files, and Pandas DataFrames. This is useful for specifying the input and expected output of test cases. By doing it this way, the data is in source control, which makes it easier to keep the test data in sync with the test cases.

col1,col2,col3,col4
1,a,true,1.23
2,b,false,4.56

from snowflake.snowpark.types import StructType, StructField, IntegerType, BooleanType, StringType, DoubleType


# Put file onto stage
session.file.put("data.csv", "@mystage", auto_compress=False)
schema = StructType(
    [
        StructField("col1", IntegerType()),
        StructField("col2", StringType()),
        StructField("col3", BooleanType()),
        StructField("col4", DoubleType()),
    ]
)

# with option SKIP_HEADER set to 1, the header will be skipped when the csv file is loaded
dataframe = session.read.schema(schema).option("SKIP_HEADER", 1).csv("@mystage/data.csv")
dataframe.show()

-------------------------------------
|"COL1"  |"COL2"  |"COL3"  |"COL4"  |
-------------------------------------
|1       |a       |True    |1.23    |
|2       |b       |False   |4.56    |
-------------------------------------

import pandas as pd

pandas_df = pd.DataFrame(
    data={
        "col1": pd.Series(["value1", "value2"]),
        "col2": pd.Series([1.23, 4.56]),
        "col3": pd.Series([123, 456]),
        "col4": pd.Series([True, False]),
    }
)

dataframe = session.create_dataframe(data=pandas_df)
dataframe.show()

-------------------------------------
|"col1"  |"col2"  |"col3"  |"col4"  |
-------------------------------------
|value1  |1.23    |123     |True    |
|value2  |4.56    |456     |False   |
-------------------------------------

from snowflake.snowpark.types import StructType, StructField, StringType, DoubleType, LongType, BooleanType

dataframe = session.create_dataframe(
    data=[
        ["value1", 1.23, 123, True],
        ["value2", 4.56, 456, False],
    ],
    schema=StructType([
        StructField("col1", StringType()),
        StructField("col2", DoubleType()),
        StructField("col3", LongType()),
        StructField("col4", BooleanType()),
    ])
)

pandas_dataframe = dataframe.to_pandas()
print(pandas_dataframe.to_string())

    COL1  COL2  COL3   COL4
0  value1  1.23   123   True
1  value2  4.56   456  False

Creating a PyTest Fixture for Session¶

PyTest fixtures are functions that are executed before a test (or module of tests), typically to provide data or connections to tests. In this case, create a fixture which returns a Snowpark Session object. First, create a test directory if you do not already have one. Then, in the test directory, create a file conftest.py with the following contents, where connection_parameters is a dictionary with your Snowflake account credentials. For more information about the dictionary format, see Creating a Session.

# test/conftest.py
import pytest
from snowflake.snowpark.session import Session

def pytest_addoption(parser):
    parser.addoption("--snowflake-session", action="store", default="live")

@pytest.fixture(scope='module')
def session(request) -> Session:
    if request.config.getoption('--snowflake-session') == 'local':
        return Session.builder.config('local_testing', True).create()
    else:
        return Session.builder.configs(CONNECTION_PARAMETERS).create()

The call to pytest_addoption adds a command line option named snowflake-session to the pytest command. The Session fixture checks this command line option, and creates a local or live Session depending on its value. This allows you to easily switch between local and live modes for testing.

# Using local mode:
pytest --snowflake-session local

# Using live mode
pytest

SQL Operations¶

Session.sql(...) is not supported in the local testing framework. Use Snowpark’s DataFrame APIs whenever possible, and in cases where you must use Session.sql(...), you can mock the tabular return value using Python’s unittest.mock.patch to patch the expected response from a given Session.sql() call.

In the example below, mock_sql() maps the SQL query text to the desired DataFrame response. The following conditional statement checks if the current session is using local testing, and if so, applies the patch to the Session.sql() method.

from unittest import mock
from functools import partial

def test_something(pytestconfig, session):

    def mock_sql(session, sql_string):  # patch for SQL operations
        if sql_string == "select 1,2,3":
            return session.create_dataframe([[1,2,3]])
        else:
            raise RuntimeError(f"Unexpected query execution: {sql_string}")

    if pytestconfig.getoption('--snowflake-session') == 'local':
        with mock.patch.object(session, 'sql', wraps=partial(mock_sql, session)): # apply patch for SQL operations
            assert session.sql("select 1,2,3").collect() == [Row(1,2,3)]
    else:
        assert session.sql("select 1,2,3").collect() == [Row(1,2,3)]

When local testing is enabled, all tables created by DataFrame.save_as_table() are saved as temporary tables in memory and can be retrieved using Session.table(). You can use the supported DataFrame operations on the table as usual.

Patching Built-In Functions¶

Not all of the built-in functions under snowflake.snowpark.functions are supported in the local testing framework. If you use a function that is not supported, you need to use the @patch decorator from snowflake.snowpark.mock to create a patch.

To define and implement the patched function, the signature (parameter list) must align with the built-in function’s parameters. The local testing framework passes parameters to the patched function using the following rules:

• For parameters of type ColumnOrName in the signature of built-in functions, ColumnEmulator is passed as the parameter of the patched functions. ColumnEmulator is similar to a pandas.Series object that contains the column data.

• For parameters of type LiteralType in the signature of built-in functions, the literal value is passed as the parameter of the patched functions.

• Otherwise, the raw value is passed as the parameter of the patched functions.

As for the returning type of the patched functions, returning an instance of ColumnEmulator is expected in correspondence with the returning type of Column of built-in functions.

For example, the built-in function to_timestamp() could be patched like this:

import datetime
from snowflake.snowpark.mock import patch, ColumnEmulator, ColumnType
from snowflake.snowpark.functions import to_timestamp
from snowflake.snowpark.types import TimestampType

@patch(to_timestamp)
def mock_to_timestamp(column: ColumnEmulator, format = None) -> ColumnEmulator:
    ret_column = ColumnEmulator(data=[datetime.datetime.strptime(row, '%Y-%m-%dT%H:%M:%S%z') for row in column])
    ret_column.sf_type = ColumnType(TimestampType(), True)
    return ret_column

Skipping Test Cases¶

If your PyTest test suite contains a test case that is not well supported by local testing, you can skip those cases using PyTest’s mark.skipif decorator. The example below assumes that you configured your session and parameters as described earlier. The condition checks if the local_testing_mode is set to local, and if so, the test case is skipped with a message explaining why it was skipped.

import pytest

@pytest.mark.skipif(
    condition="config.getvalue('local_testing_mode') == 'local'",
reason="Test case disabled for local testing"
)
def test_case(session):
    ...

Limitations¶

The following functionality is not supported:

• Raw SQL strings and operations that require parsing SQL strings. For example, session.sql and DataFrame.filter("col1 > 12") are not supported.

• UDFs, UDTFs, and stored procedures.

• Table functions.

• AsyncJobs.

• Session operations, like altering warehouses, schemas, and other session properties.

• Geometry and Geography data types.

• Aggregating window functions.

  # Selecting window function expressions is supported
  df.select("key", "value", sum_("value").over(), avg("value").over())
  
  # Aggregating window function expressions is NOT supported
  df.group_by("key").agg([sum_("value"), sum_(sum_("value")).over(window) - sum_("value")])
  

  Copy

Other limitations are:

• Variant, Array, and Object data types are only supported with standard JSON encoding and decoding. Expressions like [1,2,,3,] are considered valid JSON in Snowflake but not in local testing, where Python’s built-in JSON functionalities are used. You can specify the module-level variables snowflake.snowpark.mock.CUSTOM_JSON_ENCODER and snowflake.snowpark.mock.CUSTOM_JSON_DECODER to override the default settings.

• Only a subset of Snowflake’s functions (including window functions) are implemented. See Patching Built-In Functions to learn how to inject your own function definition.

• Patching rank related functions is currently not supported.

• Selecting columns with the same name will only return one column. As a workaround, rename the columns to have distinct names using Column.alias.

  df.select(lit(1), lit(1)).show() # col("a"), col("a")
  #---------
  #|"'1'"  |
  #---------
  #|1      |
  #|...    |
  #---------
  
  # Workaround: Column.alias
  DataFrame.select(lit(1).alias("col1_1"), lit(1).alias("col1_2"))
  # "col1_1", "col1_2"
  

  Copy

• Explicit type casting using Column.cast has the limitation that format strings are not supported for inputs: to_decimal, to_number, to_numeric, to_double, to_date, to_time, to_timestamp and outputs: to_char, to_varchar, to_binary.

• JSON strings stored in VariantType cannot be converted to Datetime types.

• For Table.merge and Table.update the implementation only supports behavior when the session parameters ERROR_ON_NONDETERMINISTIC_UPDATE and ERROR_ON_NONDETERMINISTIC_MERGE are set to False. This means for multi-joins, it updates one of the matched rows.

List of Supported APIs¶