Unit Test and create_test not working for Snowpark models

[zunedz]

I have Python Snowpark models and I'm trying to write unit tests. Below is the model definition and yaml test file:

Model definition:

@model(
    name="model_name",
    kind={
        "name": ModelKindName.SCD_TYPE_2_BY_TIME,
        "unique_key": [
             ...
        ],
        "disable_restatement": False,
        "on_destructive_change": "allow",
    },
    start="1982-01-01",
    columns={
        ...
    },
)
def execute(
    context: ExecutionContext,
    start: datetime,
    end: datetime,
    **kwargs,
) -> DataFrame:
    ...
    table = context.snowpark.table("table_name")
    ...

Test YAML Configuration:

test_model_name_example:
  gateway: gateway
  model: model_name
  inputs:
    input:
      format: csv
      path: input_path
  outputs:
    partial: true
    query:
      rows: ...

When I run sqlmesh test or sqlmesh create_test I get this error message:

ERROR: Traceback (most recent call last):
  File "/usr/local/lib/python3.11/site-packages/sqlmesh/core/model/definition.py", line 1769, in render
    df_or_iter = env[self.entrypoint](
                 ^^^^^^^^^^^^^^^^^^^^^
  File "<string>", line 13, in execute
AttributeError: 'NoneType' object has no attribute 'table'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/usr/local/lib/python3.11/site-packages/sqlmesh/core/test/definition.py", line 714, in runTest
    actual_df = self._execute_model()
                ^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/sqlmesh/core/test/definition.py", line 725, in _execute_model
    df = next(self.model.render(context=self.context, **time_kwargs, **variables))
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/sqlmesh/core/model/definition.py", line 1780, in render
    raise PythonModelEvalError(format_evaluated_code_exception(e, self.python_env))
sqlmesh.utils.errors.PythonModelEvalError:   File '' (or imported file), line 13, in execute
    def execute(context: ExecutionContext, start: datetime, end: datetime, **kwargs
        ):
        table = context.snowpark.table(
      AttributeError: 'NoneType' object has no attribute 'table'

----------------------------------------------------------------------