Repository: TobikoData/sqlmesh
Pull request #2435 (2024-04-13 00:42:35)

Feat!: allow overriding input fixture schema in unit tests
This PR:

- Adds a new YAML setting, `columns`, which allows users to define schemas for upstream models ([context](https://tobiko-data.slack.com/archives/C044BRE5W4S/p1712872939106209))
- Refactors several pieces of the test definition module:
  - Moves `TestError` to `sqlmesh.utils.errors` where it belongs (breaking)
  - Moves most validations inside of `_validate_and_normalize_test` (previously `_normalize_test`)
  - Introduces a couple new in-memory caches to avoid repeating the same normalization calls


Patch:
diff --git a/docs/guides/tablediff.md b/docs/guides/tablediff.md
index ff2bc144..f3550aab 100644
--- a/docs/guides/tablediff.md
+++ b/docs/guides/tablediff.md
@@ -29,21 +29,21 @@ For example, we could make two modifications to the [SQLMesh quickstart](../quic
 MODEL (
     name sqlmesh_example.incremental_model,
     kind INCREMENTAL_BY_TIME_RANGE (
-        time_column ds
+        time_column event_date
     ),
     start '2020-01-01',
     cron '@daily',
-    grain (id, ds)
+    grain (id, event_date)
 );
 
 SELECT
     id,
     CASE WHEN item_id = 3 THEN 4 ELSE item_id END as item_id, -- Change item_id 3 to 4
-    ds,
+    event_date,
 FROM
     sqlmesh_example.seed_model
 WHERE
-    ds between @start_ds and @end_ds
+    event_date between @start_ds and @end_ds
     AND id != 1 -- Remove row whose item_id is 1
 ```
 
diff --git a/sqlmesh/core/engine_adapter/base.py b/sqlmesh/core/engine_adapter/base.py
index 814e5876..feb1d09a 100644
--- a/sqlmesh/core/engine_adapter/base.py
+++ b/sqlmesh/core/engine_adapter/base.py
@@ -946,7 +946,7 @@ class EngineAdapter:
         view_name: TableName,
         ignore_if_not_exists: bool = True,
         materialized: bool = False,
-        cascade: bool = False,
+        **kwargs: t.Any,
     ) -> None:
         """Drop a view."""
         self.execute(
@@ -954,8 +954,8 @@ class EngineAdapter:
                 this=exp.to_table(view_name),
                 exists=ignore_if_not_exists,
                 materialized=materialized and self.SUPPORTS_MATERIALIZED_VIEWS,
-                cascade=cascade,
                 kind="VIEW",
+                **kwargs,
             )
         )
 
diff --git a/sqlmesh/core/engine_adapter/base_postgres.py b/sqlmesh/core/engine_adapter/base_postgres.py
index 049f84ee..4ca5627f 100644
--- a/sqlmesh/core/engine_adapter/base_postgres.py
+++ b/sqlmesh/core/engine_adapter/base_postgres.py
@@ -93,7 +93,7 @@ class BasePostgresEngineAdapter(EngineAdapter):
         """
         with self.transaction():
             if replace:
-                self.drop_view(view_name, materialized=materialized, cascade=True)
+                self.drop_view(view_name, materialized=materialized)
             super().create_view(
                 view_name,
                 query_or_df,
@@ -105,6 +105,21 @@ class BasePostgresEngineAdapter(EngineAdapter):
                 **create_kwargs,
             )
 
+    def drop_view(
+        self,
+        view_name: TableName,
+        ignore_if_not_exists: bool = True,
+        materialized: bool = False,
+        **kwargs: t.Any,
+    ) -> None:
+        kwargs["cascade"] = kwargs.get("cascade", True)
+        return super().drop_view(
+            view_name,
+            ignore_if_not_exists=ignore_if_not_exists,
+            materialized=materialized,
+            **kwargs,
+        )
+
     def _get_data_objects(
         self, schema_name: SchemaName, object_names: t.Optional[t.Set[str]] = None
     ) -> t.List[DataObject]:
diff --git a/sqlmesh/core/snapshot/evaluator.py b/sqlmesh/core/snapshot/evaluator.py
index ec036d6d..0e68ed05 100644
--- a/sqlmesh/core/snapshot/evaluator.py
+++ b/sqlmesh/core/snapshot/evaluator.py
@@ -1047,7 +1047,7 @@ class EmbeddedStrategy(SymbolicStrategy):
     ) -> None:
         target_name = view_name.for_environment(environment_naming_info)
         logger.info("Dropping view '%s' for non-materialized table", target_name)
-        self.adapter.drop_view(target_name)
+        self.adapter.drop_view(target_name, cascade=False)
 
 
 class PromotableStrategy(EvaluationStrategy):
@@ -1075,7 +1075,7 @@ class PromotableStrategy(EvaluationStrategy):
     ) -> None:
         target_name = view_name.for_environment(environment_naming_info)
         logger.info("Dropping view '%s'", target_name)
-        self.adapter.drop_view(target_name)
+        self.adapter.drop_view(target_name, cascade=False)
 
 
 class MaterializableStrategy(PromotableStrategy):
diff --git a/sqlmesh/utils/errors.py b/sqlmesh/utils/errors.py
index dc053fa1..21cb857f 100644
--- a/sqlmesh/utils/errors.py
+++ b/sqlmesh/utils/errors.py
@@ -100,6 +100,10 @@ class AuditError(SQLMeshError):
         return self.query.sql(dialect=dialect or self.adapter_dialect, **opts)
 
 
+class TestError(SQLMeshError):
+    pass
+
+
 class NotificationTargetError(SQLMeshError):
     pass
 


Test patch:
diff --git a/sqlmesh/core/test/definition.py b/sqlmesh/core/test/definition.py
index 5b28201c..d6443023 100644
--- a/sqlmesh/core/test/definition.py
+++ b/sqlmesh/core/test/definition.py
@@ -20,15 +20,11 @@ from sqlmesh.core.engine_adapter import EngineAdapter
 from sqlmesh.core.model import Model, PythonModel, SqlModel
 from sqlmesh.utils import UniqueKeyDict, random_id, yaml
 from sqlmesh.utils.date import pandas_timestamp_to_pydatetime
-from sqlmesh.utils.errors import ConfigError, SQLMeshError
+from sqlmesh.utils.errors import ConfigError, TestError
 
 Row = t.Dict[str, t.Any]
 
 
-class TestError(SQLMeshError):
-    """Test error"""
-
-
 class ModelTest(unittest.TestCase):
     __test__ = False
 
@@ -66,19 +62,19 @@ class ModelTest(unittest.TestCase):
         self.default_catalog = default_catalog
         self.dialect = dialect
 
-        self._normalize_test(dialect)
+        self._fixture_table_cache: t.Dict[str, exp.Table] = {}
+        self._normalized_column_name_cache: t.Dict[str, str] = {}
+        self._normalized_model_name_cache: t.Dict[t.Tuple[str, bool], str] = {}
 
-        inputs = [
-            normalize_model_name(input, default_catalog=default_catalog, dialect=dialect)
-            for input in self.body.get("inputs", {})
-        ]
-        for depends_on in self.model.depends_on:
-            if depends_on not in inputs:
-                _raise_error(f"Incomplete test, missing input for table {depends_on}", path)
+        self._validate_and_normalize_test()
+
+        # This ID is appended to each input fixture name to avoid concurrency issues
+        self._test_id = random_id(short=True)
 
         self._engine_adapter_dialect = Dialect.get_or_raise(self.engine_adapter.dialect)
         self._transforms = self._engine_adapter_dialect.generator_class.TRANSFORMS
 
+        # When execution_time is set, we mock the CURRENT_* SQL expressions so they always return it
         self._execution_time = str(self.body.get("vars", {}).get("execution_time") or "")
         if self._execution_time:
             exec_time = exp.Literal.string(self._execution_time)
@@ -90,9 +86,6 @@ class ModelTest(unittest.TestCase):
                 exp.CurrentTimestamp: lambda self, _: self.sql(exp.cast(exec_time, "timestamp")),
             }
 
-        self._test_id = random_id(short=True)
-        self._fixture_table_cache: t.Dict[str, exp.Table] = {}
-
         super().__init__()
 
     def shortDescription(self) -> t.Optional[str]:
@@ -100,11 +93,13 @@ class ModelTest(unittest.TestCase):
 
     def setUp(self) -> None:
         """Load all input tables"""
-        for table_name, values in self.body.get("inputs", {}).items():
-            columns_to_types: t.Dict[str, exp.DataType] = {}
-            if table_name in self.models:
-                model = self.models[table_name]
-                columns_to_types = model.columns_to_types if model.annotated else {}
+        for name, values in self.body.get("inputs", {}).items():
+            # Types specified in the test take precedence over the corresponding inferred ones
+            model = self.models.get(name)
+            columns_to_types = {
+                **(model.columns_to_types if model and model.annotated else {}),
+                **values.get("columns", {}),
+            }
 
             rows = values["rows"]
             if not columns_to_types and rows:
@@ -114,7 +109,7 @@ class ModelTest(unittest.TestCase):
                         v_type, into=exp.DataType, dialect=self.dialect
                     )
 
-            test_fixture_table = self._test_fixture_table(table_name)
+            test_fixture_table = self._test_fixture_table(name)
             if test_fixture_table.db:
                 self.engine_adapter.create_schema(
                     schema_(test_fixture_table.args["db"], test_fixture_table.args.get("catalog"))
@@ -126,8 +121,8 @@ class ModelTest(unittest.TestCase):
     def tearDown(self) -> None:
         """Drop all fixture tables."""
         if not self.preserve_fixtures:
-            for table_name in self.body.get("inputs", {}):
-                self.engine_adapter.drop_view(self._fixture_table_cache[table_name])
+            for name in self.body.get("inputs", {}):
+                self.engine_adapter.drop_view(self._fixture_table_cache[name])
 
     def assert_equal(
         self,
@@ -221,64 +216,55 @@ class ModelTest(unittest.TestCase):
             path: An optional path to the test definition yaml file.
             preserve_fixtures: Preserve the fixture tables in the testing database, useful for debugging.
         """
-        if "model" not in body:
-            _raise_error("Incomplete test, missing model name", path)
+        name = normalize_model_name(body["model"], default_catalog=default_catalog, dialect=dialect)
+        model = models.get(name)
+        if not model:
+            _raise_error(f"Model '{name}' was not found", path)
 
-        if "outputs" not in body:
-            _raise_error("Incomplete test, missing outputs", path)
-
-        model_name = normalize_model_name(
-            body["model"], default_catalog=default_catalog, dialect=dialect
-        )
-        if model_name not in models:
-            _raise_error(f"Model '{model_name}' was not found", path)
-
-        model = models[model_name]
         if isinstance(model, SqlModel):
-            return SqlModelTest(
-                body,
-                test_name,
-                model,
-                models,
-                engine_adapter,
-                dialect,
-                path,
-                preserve_fixtures,
-                default_catalog,
-            )
-        if isinstance(model, PythonModel):
-            return PythonModelTest(
-                body,
-                test_name,
-                model,
-                models,
-                engine_adapter,
-                dialect,
-                path,
-                preserve_fixtures,
-                default_catalog,
-            )
+            test_type: t.Type[ModelTest] = SqlModelTest
+        elif isinstance(model, PythonModel):
+            test_type = PythonModelTest
+        else:
+            _raise_error(f"Model '{name}' is an unsupported model type for testing", path)
 
-        raise TestError(f"Model '{model_name}' is an unsupported model type for testing at {path}")
+        return test_type(
+            body,
+            test_name,
+            t.cast(Model, model),
+            models,
+            engine_adapter,
+            dialect,
+            path,
+            preserve_fixtures,
+            default_catalog,
+        )
 
     def __str__(self) -> str:
         return f"{self.test_name} ({self.path})"
 
-    def _normalize_test(self, dialect: str | None) -> None:
-        """Normalizes all identifiers in this test according to the given dialect."""
+    def _validate_and_normalize_test(self) -> None:
+        inputs = self.body.get("inputs")
+        outputs = self.body.get("outputs", {})
+
+        if not outputs:
+            _raise_error("Incomplete test, missing outputs", self.path)
 
-        def _normalize_rows(values: t.List[Row] | t.Dict, partial: bool = False) -> t.Dict:
+        ctes = outputs.get("ctes")
+        query = outputs.get("query")
+        partial = outputs.pop("partial", None)
+
+        def _normalize_rows(
+            values: t.List[Row] | t.Dict, name: str, partial: bool = False
+        ) -> t.Dict:
             if not isinstance(values, dict):
                 values = {"rows": values}
 
             if "rows" not in values:
-                _raise_error("Incomplete test, missing row data for table", self.path)
+                _raise_error(f"Incomplete test, missing row data for '{name}'", self.path)
 
             values["rows"] = [
-                {
-                    normalize_identifiers(column, dialect=dialect).name: value
-                    for column, value in row.items()
-                }
+                {self._normalize_column_name(column): value for column, value in row.items()}
                 for row in values["rows"]
             ]
             if partial:
@@ -286,30 +272,47 @@ class ModelTest(unittest.TestCase):
 
             return values
 
-        def _normalize_sources(sources: t.Dict, partial: bool = False) -> t.Dict:
+        def _normalize_sources(
+            sources: t.Dict, partial: bool = False, with_default_catalog: bool = True
+        ) -> t.Dict:
             return {
-                normalize_model_name(
-                    name, default_catalog=self.default_catalog, dialect=dialect
-                ): _normalize_rows(values, partial=partial)
+                self._normalize_model_name(
+                    name, with_default_catalog=with_default_catalog
+                ): _normalize_rows(values, name, partial=partial)
                 for name, values in sources.items()
             }
 
-        inputs = self.body.get("inputs")
-        outputs = self.body["outputs"]
-        ctes = outputs.get("ctes")
-        query = outputs.get("query")
-        partial = outputs.pop("partial", None)
-
         if inputs:
-            self.body["inputs"] = _normalize_sources(inputs)
+            inputs = _normalize_sources(inputs)
+            for name, values in inputs.items():
+                columns = values.get("columns")
+                if columns is None:
+                    continue
+
+                if not isinstance(columns, dict):
+                    _raise_error(
+                        f"Invalid 'columns' value for model '{name}', expected a mapping name -> type",
+                        self.path,
+                    )
+
+                values["columns"] = {
+                    self._normalize_column_name(c): exp.DataType.build(t, dialect=self.dialect)
+                    for c, t in columns.items()
+                }
+
+            for depends_on in self.model.depends_on:
+                if depends_on not in inputs:
+                    _raise_error(f"Incomplete test, missing input model '{depends_on}'", self.path)
+
+            self.body["inputs"] = inputs
+
         if ctes:
-            outputs["ctes"] = _normalize_sources(ctes, partial=partial)
+            outputs["ctes"] = _normalize_sources(ctes, partial=partial, with_default_catalog=False)
+
         if query or query == []:
-            outputs["query"] = _normalize_rows(query, partial=partial)
+            outputs["query"] = _normalize_rows(query, self.model.name, partial=partial)
 
-        self.body["model"] = normalize_model_name(
-            self.body["model"], default_catalog=self.default_catalog, dialect=dialect
-        )
+        self.body["model"] = self._normalize_model_name(self.body["model"])
 
     def _test_fixture_table(self, name: str) -> exp.Table:
         table = self._fixture_table_cache.get(name)
@@ -320,6 +323,25 @@ class ModelTest(unittest.TestCase):
 
         return table
 
+    def _normalize_model_name(self, name: str, with_default_catalog: bool = True) -> str:
+        normalized_name = self._normalized_model_name_cache.get((name, with_default_catalog))
+        if normalized_name is None:
+            default_catalog = self.default_catalog if with_default_catalog else None
+            normalized_name = normalize_model_name(
+                name, default_catalog=default_catalog, dialect=self.dialect
+            )
+            self._normalized_model_name_cache[(name, with_default_catalog)] = normalized_name
+
+        return normalized_name
+
+    def _normalize_column_name(self, name: str) -> str:
+        normalized_name = self._normalized_column_name_cache.get(name)
+        if normalized_name is None:
+            normalized_name = normalize_identifiers(name, dialect=self.dialect).name
+            self._normalized_column_name_cache[name] = normalized_name
+
+        return normalized_name
+
 
 class SqlModelTest(ModelTest):
     def test_ctes(self, ctes: t.Dict[str, exp.Expression]) -> None:
@@ -349,7 +371,7 @@ class SqlModelTest(ModelTest):
         mapping = {
             name: self._test_fixture_table(name).sql()
             for name in [
-                normalize_model_name(name, self.default_catalog, self.dialect)
+                self._normalize_model_name(name)
                 for name in self.models.keys() | self.body.get("inputs", {}).keys()
             ]
         }
@@ -360,7 +382,10 @@ class SqlModelTest(ModelTest):
         )
 
         self.test_ctes(
-            {normalize_model_name(cte.alias, None, self.dialect): cte for cte in query.ctes}
+            {
+                self._normalize_model_name(cte.alias, with_default_catalog=False): cte
+                for cte in query.ctes
+            }
         )
 
         values = self.body["outputs"].get("query")
@@ -385,7 +410,7 @@ class PythonModelTest(ModelTest):
         self,
         body: t.Dict[str, t.Any],
         test_name: str,
-        model: PythonModel,
+        model: Model,
         models: UniqueKeyDict[str, Model],
         engine_adapter: EngineAdapter,
         dialect: str | None = None,
diff --git a/tests/core/engine_adapter/test_base_postgres.py b/tests/core/engine_adapter/test_base_postgres.py
index 6b1258c4..a7a9369f 100644
--- a/tests/core/engine_adapter/test_base_postgres.py
+++ b/tests/core/engine_adapter/test_base_postgres.py
@@ -50,3 +50,23 @@ def test_create_view(make_mocked_engine_adapter: t.Callable):
             call('CREATE VIEW "db"."view" AS SELECT 1'),
         ]
     )
+
+
+def test_drop_view(make_mocked_engine_adapter: t.Callable):
+    adapter = make_mocked_engine_adapter(BasePostgresEngineAdapter)
+    adapter.SUPPORTS_MATERIALIZED_VIEWS = True
+
+    adapter.drop_view("db.view")
+    adapter.drop_view("db.view", materialized=True)
+    adapter.drop_view("db.view", cascade=False)
+
+    adapter.cursor.execute.assert_has_calls(
+        [
+            # 1st call
+            call('DROP VIEW IF EXISTS "db"."view" CASCADE'),
+            # 2nd call
+            call('DROP MATERIALIZED VIEW IF EXISTS "db"."view" CASCADE'),
+            # 3rd call
+            call('DROP VIEW IF EXISTS "db"."view"'),
+        ]
+    )
diff --git a/tests/core/test_snapshot_evaluator.py b/tests/core/test_snapshot_evaluator.py
index 6680f105..64b1cff2 100644
--- a/tests/core/test_snapshot_evaluator.py
+++ b/tests/core/test_snapshot_evaluator.py
@@ -261,6 +261,26 @@ def test_promote(mocker: MockerFixture, adapter_mock, make_snapshot):
     )
 
 
+def test_demote(mocker: MockerFixture, adapter_mock, make_snapshot):
+    evaluator = SnapshotEvaluator(adapter_mock)
+
+    model = SqlModel(
+        name="test_schema.test_model",
+        kind=ViewKind(),
+        query=parse_one("SELECT a FROM tbl"),
+    )
+
+    snapshot = make_snapshot(model)
+    snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
+
+    evaluator.demote([snapshot], EnvironmentNamingInfo(name="test_env"))
+
+    adapter_mock.drop_view.assert_called_once_with(
+        "test_schema__test_env.test_model",
+        cascade=False,
+    )
+
+
 def test_promote_default_catalog(adapter_mock, make_snapshot):
     evaluator = SnapshotEvaluator(adapter_mock)
 
diff --git a/tests/core/test_test.py b/tests/core/test_test.py
index 6f357236..216a27f3 100644
--- a/tests/core/test_test.py
+++ b/tests/core/test_test.py
@@ -14,8 +14,8 @@ from sqlmesh.core import constants as c
 from sqlmesh.core.config import Config, DuckDBConnectionConfig, ModelDefaultsConfig
 from sqlmesh.core.context import Context
 from sqlmesh.core.dialect import parse
-from sqlmesh.core.model import PythonModel, SqlModel, load_sql_based_model, model
-from sqlmesh.core.test.definition import PythonModelTest, SqlModelTest
+from sqlmesh.core.model import Model, SqlModel, load_sql_based_model, model
+from sqlmesh.core.test.definition import ModelTest, PythonModelTest, SqlModelTest
 from sqlmesh.utils.errors import ConfigError
 from sqlmesh.utils.yaml import load as load_yaml
 
@@ -27,19 +27,9 @@ pytestmark = pytest.mark.slow
 SUSHI_FOO_META = "MODEL (name sushi.foo, kind FULL)"
 
 
-@t.overload
 def _create_test(
-    body: t.Dict[str, t.Any], test_name: str, model: SqlModel, context: Context
-) -> SqlModelTest: ...
-
-
-@t.overload
-def _create_test(
-    body: t.Dict[str, t.Any], test_name: str, model: PythonModel, context: Context
-) -> PythonModelTest: ...
-
-
-def _create_test(body, test_name, model, context):
+    body: t.Dict[str, t.Any], test_name: str, model: Model, context: Context
+) -> ModelTest:
     test_type = SqlModelTest if isinstance(model, SqlModel) else PythonModelTest
     return test_type(
         body=body[test_name],
@@ -114,9 +104,9 @@ def full_model_with_two_ctes(request) -> SqlModel:
 
 
 def test_ctes(sushi_context: Context, full_model_with_two_ctes: SqlModel) -> None:
-    model = t.cast(SqlModel, sushi_context.upsert_model(full_model_with_two_ctes))
-    body = load_yaml(
-        """
+    test = _create_test(
+        body=load_yaml(
+            """
 test_foo:
   model: sushi.foo
   inputs:
@@ -133,9 +123,12 @@ test_foo:
   vars:
     start: 2022-01-01
     end: 2022-01-01
-        """
+            """
+        ),
+        test_name="test_foo",
+        model=sushi_context.upsert_model(full_model_with_two_ctes),
+        context=sushi_context,
     )
-    test = _create_test(body, "test_foo", model, sushi_context)
 
     random_id = "jzngz56a"
     test._test_id = random_id
@@ -147,9 +140,10 @@ test_foo:
 
 
 def test_ctes_only(sushi_context: Context, full_model_with_two_ctes: SqlModel) -> None:
-    model = t.cast(SqlModel, sushi_context.upsert_model(full_model_with_two_ctes))
-    body = load_yaml(
-        """
+    _check_successful_or_raise(
+        _create_test(
+            body=load_yaml(
+                """
 test_foo:
   model: sushi.foo
   inputs:
@@ -164,16 +158,20 @@ test_foo:
   vars:
     start: 2022-01-01
     end: 2022-01-01
-        """
+                """
+            ),
+            test_name="test_foo",
+            model=sushi_context.upsert_model(full_model_with_two_ctes),
+            context=sushi_context,
+        ).run()
     )
-    result = _create_test(body, "test_foo", model, sushi_context).run()
-    _check_successful_or_raise(result)
 
 
 def test_query_only(sushi_context: Context, full_model_with_two_ctes: SqlModel) -> None:
-    model = t.cast(SqlModel, sushi_context.upsert_model(full_model_with_two_ctes))
-    body = load_yaml(
-        """
+    _check_successful_or_raise(
+        _create_test(
+            body=load_yaml(
+                """
 test_foo:
   model: sushi.foo
   inputs:
@@ -185,16 +183,20 @@ test_foo:
   vars:
     start: 2022-01-01
     end: 2022-01-01
-        """
+                """
+            ),
+            test_name="test_foo",
+            model=sushi_context.upsert_model(full_model_with_two_ctes),
+            context=sushi_context,
+        ).run()
     )
-    result = _create_test(body, "test_foo", model, sushi_context).run()
-    _check_successful_or_raise(result)
 
 
 def test_with_rows(sushi_context: Context, full_model_with_single_cte: SqlModel) -> None:
-    model = t.cast(SqlModel, sushi_context.upsert_model(full_model_with_single_cte))
-    body = load_yaml(
-        """
+    _check_successful_or_raise(
+        _create_test(
+            body=load_yaml(
+                """
 test_foo:
   model: sushi.foo
   inputs:
@@ -212,16 +214,20 @@ test_foo:
   vars:
     start: 2022-01-01
     end: 2022-01-01
-        """
+                """
+            ),
+            test_name="test_foo",
+            model=sushi_context.upsert_model(full_model_with_single_cte),
+            context=sushi_context,
+        ).run()
     )
-    result = _create_test(body, "test_foo", model, sushi_context).run()
-    _check_successful_or_raise(result)
 
 
 def test_without_rows(sushi_context: Context, full_model_with_single_cte: SqlModel) -> None:
-    model = t.cast(SqlModel, sushi_context.upsert_model(full_model_with_single_cte))
-    body = load_yaml(
-        """
+    _check_successful_or_raise(
+        _create_test(
+            body=load_yaml(
+                """
 test_foo:
   model: sushi.foo
   inputs:
@@ -236,16 +242,20 @@ test_foo:
   vars:
     start: 2022-01-01
     end: 2022-01-01
-        """
+                """
+            ),
+            test_name="test_foo",
+            model=sushi_context.upsert_model(full_model_with_single_cte),
+            context=sushi_context,
+        ).run()
     )
-    result = _create_test(body, "test_foo", model, sushi_context).run()
-    _check_successful_or_raise(result)
 
 
 def test_column_order(sushi_context: Context, full_model_without_ctes: SqlModel) -> None:
-    model = t.cast(SqlModel, sushi_context.upsert_model(full_model_without_ctes))
-    body = load_yaml(
-        """
+    _check_successful_or_raise(
+        _create_test(
+            body=load_yaml(
+                """
 test_foo:
   model: sushi.foo
   inputs:
@@ -261,10 +271,13 @@ test_foo:
   vars:
     start: 2022-01-01
     end: 2022-01-01
-        """
+                """
+            ),
+            test_name="test_foo",
+            model=sushi_context.upsert_model(full_model_without_ctes),
+            context=sushi_context,
+        ).run()
     )
-    result = _create_test(body, "test_foo", model, sushi_context).run()
-    _check_successful_or_raise(result)
 
 
 def test_row_order(sushi_context: Context, full_model_without_ctes: SqlModel) -> None:
@@ -320,14 +333,10 @@ test_foo:
 
 
 def test_partial_data(sushi_context: Context) -> None:
-    model = _create_model(
-        "WITH source AS (SELECT id, name FROM sushi.waiter_names) SELECT id, name, 'nan' as str FROM source",
-        default_catalog=sushi_context.default_catalog,
-    )
-    model = t.cast(SqlModel, sushi_context.upsert_model(model))
-
-    body = load_yaml(
-        """
+    _check_successful_or_raise(
+        _create_test(
+            body=load_yaml(
+                """
 test_foo:
   model: sushi.foo
   inputs:
@@ -353,16 +362,26 @@ test_foo:
       - id: 3
         name: 'bob'
         str: nan
-        """
+                """
+            ),
+            test_name="test_foo",
+            model=sushi_context.upsert_model(
+                _create_model(
+                    "WITH source AS (SELECT id, name FROM sushi.waiter_names) "
+                    "SELECT id, name, 'nan' as str FROM source",
+                    default_catalog=sushi_context.default_catalog,
+                )
+            ),
+            context=sushi_context,
+        ).run()
     )
-    result = _create_test(body, "test_foo", model, sushi_context).run()
-    _check_successful_or_raise(result)
 
 
 def test_partial_output_columns() -> None:
-    result = _create_test(
-        body=load_yaml(
-            """
+    _check_successful_or_raise(
+        _create_test(
+            body=load_yaml(
+                """
 test_foo:
   model: sushi.foo
   inputs:
@@ -387,18 +406,18 @@ test_foo:
           b: 2
         - a: 5
           b: 6
-            """
-        ),
-        test_name="test_foo",
-        model=_create_model("WITH t AS (SELECT a, b, c, d FROM raw) SELECT a, b, c, d FROM t"),
-        context=Context(config=Config(model_defaults=ModelDefaultsConfig(dialect="duckdb"))),
-    ).run()
-
-    _check_successful_or_raise(result)
+                """
+            ),
+            test_name="test_foo",
+            model=_create_model("WITH t AS (SELECT a, b, c, d FROM raw) SELECT a, b, c, d FROM t"),
+            context=Context(config=Config(model_defaults=ModelDefaultsConfig(dialect="duckdb"))),
+        ).run()
+    )
 
-    result = _create_test(
-        body=load_yaml(
-            """
+    _check_successful_or_raise(
+        _create_test(
+            body=load_yaml(
+                """
 test_foo:
   model: sushi.foo
   inputs:
@@ -426,23 +445,20 @@ test_foo:
         - a: 5
           b: 6
           c: 7
-            """
-        ),
-        test_name="test_foo",
-        model=_create_model("WITH t AS (SELECT a, b, c, d FROM raw) SELECT a, b, c, d FROM t"),
-        context=Context(config=Config(model_defaults=ModelDefaultsConfig(dialect="duckdb"))),
-    ).run()
+                """
+            ),
+            test_name="test_foo",
+            model=_create_model("WITH t AS (SELECT a, b, c, d FROM raw) SELECT a, b, c, d FROM t"),
+            context=Context(config=Config(model_defaults=ModelDefaultsConfig(dialect="duckdb"))),
+        ).run()
+    )
 
 
 def test_partial_data_column_order(sushi_context: Context) -> None:
-    model = _create_model(
-        "SELECT id, name, price, event_date FROM sushi.items",
-        default_catalog=sushi_context.default_catalog,
-    )
-    model = t.cast(SqlModel, sushi_context.upsert_model(model))
-
-    body = load_yaml(
-        """
+    _check_successful_or_raise(
+        _create_test(
+            body=load_yaml(
+                """
 test_foo:
   model: sushi.foo
   inputs:
@@ -459,18 +475,25 @@ test_foo:
       - id: 9876
         name: hello
         event_date: 2020-01-02
-        """
+                """
+            ),
+            test_name="test_foo",
+            model=sushi_context.upsert_model(
+                _create_model(
+                    "SELECT id, name, price, event_date FROM sushi.items",
+                    default_catalog=sushi_context.default_catalog,
+                )
+            ),
+            context=sushi_context,
+        ).run()
     )
-    result = _create_test(body, "test_foo", model, sushi_context).run()
-    _check_successful_or_raise(result)
 
 
 def test_partial_data_missing_schemas(sushi_context: Context) -> None:
-    model = _create_model("SELECT * FROM unknown")
-    model = t.cast(SqlModel, sushi_context.upsert_model(model))
-
-    body = load_yaml(
-        """
+    _check_successful_or_raise(
+        _create_test(
+            body=load_yaml(
+                """
 test_foo:
   model: sushi.foo
   inputs:
@@ -483,18 +506,17 @@ test_foo:
       - a: 1
         b: bla
       - b: baz
-        """
+                """
+            ),
+            test_name="test_foo",
+            model=sushi_context.upsert_model(_create_model("SELECT * FROM unknown")),
+            context=sushi_context,
+        ).run()
     )
-    result = _create_test(body, "test_foo", model, sushi_context).run()
-    _check_successful_or_raise(result)
-
-    model = _create_model(
-        "SELECT *, DATE_TRUNC('month', date)::DATE AS month, NULL::DATE AS null_date, FROM unknown"
-    )
-    model = t.cast(SqlModel, sushi_context.upsert_model(model))
-
-    body = load_yaml(
-        """
+    _check_successful_or_raise(
+        _create_test(
+            body=load_yaml(
+                """
 test_foo:
   model: sushi.foo
   inputs:
@@ -513,16 +535,24 @@ test_foo:
         date: 2023-02-10
         month: 2023-02-01
         null_date:
-        """
+                """
+            ),
+            test_name="test_foo",
+            model=sushi_context.upsert_model(
+                _create_model(
+                    "SELECT *, DATE_TRUNC('month', date)::DATE AS month, NULL::DATE AS null_date, FROM unknown"
+                )
+            ),
+            context=sushi_context,
+        ).run()
     )
-    result = _create_test(body, "test_foo", model, sushi_context).run()
-    _check_successful_or_raise(result)
 
 
 def test_missing_column_failure(sushi_context: Context, full_model_without_ctes: SqlModel) -> None:
-    model = t.cast(SqlModel, sushi_context.upsert_model(full_model_without_ctes))
-    body = load_yaml(
-        """
+    _check_successful_or_raise(
+        _create_test(
+            body=load_yaml(
+                """
 test_foo:
   model: sushi.foo
   description: sushi.foo's output has a missing column (fails intentionally)
@@ -535,23 +565,26 @@ test_foo:
     query:
       - id: 1
         value: null
-        """
-    )
-    result = _create_test(body, "test_foo", model, sushi_context).run()
-
-    expected_failure_msg = (
-        "AssertionError: Data mismatch (exp: expected, act: actual)\n\n"
-        "  value      ds    \n"
-        "    exp act exp act\n"
-        "0   NaN   2 NaN   3\n"
+                """
+            ),
+            test_name="test_foo",
+            model=sushi_context.upsert_model(full_model_without_ctes),
+            context=sushi_context,
+        ).run(),
+        expected_msg=(
+            "AssertionError: Data mismatch (exp: expected, act: actual)\n\n"
+            "  value      ds    \n"
+            "    exp act exp act\n"
+            "0   NaN   2 NaN   3\n"
+        ),
     )
-    _check_successful_or_raise(result, expected_msg=expected_failure_msg)
 
 
 def test_row_difference_failure() -> None:
-    result = _create_test(
-        body=load_yaml(
-            """
+    _check_successful_or_raise(
+        _create_test(
+            body=load_yaml(
+                """
 test_foo:
   model: sushi.foo
   inputs:
@@ -561,24 +594,23 @@ test_foo:
     query:
       - value: 1
       - value: 2
-            """
+                """
+            ),
+            test_name="test_foo",
+            model=_create_model("SELECT value FROM raw"),
+            context=Context(config=Config(model_defaults=ModelDefaultsConfig(dialect="duckdb"))),
+        ).run(),
+        expected_msg=(
+            "AssertionError: Data mismatch (rows are different)\n\n"
+            "Missing rows:\n\n"
+            "   value\n"
+            "0      2\n"
         ),
-        test_name="test_foo",
-        model=_create_model("SELECT value FROM raw"),
-        context=Context(config=Config(model_defaults=ModelDefaultsConfig(dialect="duckdb"))),
-    ).run()
-
-    expected_error_msg = (
-        "AssertionError: Data mismatch (rows are different)\n\n"
-        "Missing rows:\n\n"
-        "   value\n"
-        "0      2\n"
     )
-    _check_successful_or_raise(result, expected_msg=expected_error_msg)
-
-    result = _create_test(
-        body=load_yaml(
-            """
+    _check_successful_or_raise(
+        _create_test(
+            body=load_yaml(
+                """
 test_foo:
   model: sushi.foo
   inputs:
@@ -587,24 +619,25 @@ test_foo:
   outputs:
     query:
       - value: 1
-            """
+                """
+            ),
+            test_name="test_foo",
+            model=_create_model(
+                "SELECT value FROM raw UNION ALL SELECT value + 1 AS value FROM raw"
+            ),
+            context=Context(config=Config(model_defaults=ModelDefaultsConfig(dialect="duckdb"))),
+        ).run(),
+        expected_msg=(
+            "AssertionError: Data mismatch (rows are different)\n\n"
+            "Unexpected rows:\n\n"
+            "   value\n"
+            "0      2\n"
         ),
-        test_name="test_foo",
-        model=_create_model("SELECT value FROM raw UNION ALL SELECT value + 1 AS value FROM raw"),
-        context=Context(config=Config(model_defaults=ModelDefaultsConfig(dialect="duckdb"))),
-    ).run()
-
-    expected_error_msg = (
-        "AssertionError: Data mismatch (rows are different)\n\n"
-        "Unexpected rows:\n\n"
-        "   value\n"
-        "0      2\n"
     )
-    _check_successful_or_raise(result, expected_msg=expected_error_msg)
-
-    result = _create_test(
-        body=load_yaml(
-            """
+    _check_successful_or_raise(
+        _create_test(
+            body=load_yaml(
+                """
 test_foo:
   model: sushi.foo
   inputs:
@@ -615,30 +648,32 @@ test_foo:
       - value: 1
       - value: 3
       - value: 4
-            """
+                """
+            ),
+            test_name="test_foo",
+            model=_create_model(
+                "SELECT value FROM raw UNION ALL SELECT value + 1 AS value FROM raw"
+            ),
+            context=Context(config=Config(model_defaults=ModelDefaultsConfig(dialect="duckdb"))),
+        ).run(),
+        expected_msg=(
+            "AssertionError: Data mismatch (rows are different)\n\n"
+            "Missing rows:\n\n"
+            "   value\n"
+            "0      3\n"
+            "1      4\n\n"
+            "Unexpected rows:\n\n"
+            "   value\n"
+            "0      2\n"
         ),
-        test_name="test_foo",
-        model=_create_model("SELECT value FROM raw UNION ALL SELECT value + 1 AS value FROM raw"),
-        context=Context(config=Config(model_defaults=ModelDefaultsConfig(dialect="duckdb"))),
-    ).run()
-
-    expected_error_msg = (
-        "AssertionError: Data mismatch (rows are different)\n\n"
-        "Missing rows:\n\n"
-        "   value\n"
-        "0      3\n"
-        "1      4\n\n"
-        "Unexpected rows:\n\n"
-        "   value\n"
-        "0      2\n"
     )
-    _check_successful_or_raise(result, expected_msg=expected_error_msg)
 
 
 def test_unknown_column_error() -> None:
-    result = _create_test(
-        body=load_yaml(
-            """
+    _check_successful_or_raise(
+        _create_test(
+            body=load_yaml(
+                """
 test_foo:
   model: sushi.foo
   inputs:
@@ -648,45 +683,49 @@ test_foo:
   outputs:
     query:
       - foo: 1
-            """
+                """
+            ),
+            test_name="test_foo",
+            model=_create_model("SELECT id, value FROM raw"),
+            context=Context(config=Config(model_defaults=ModelDefaultsConfig(dialect="duckdb"))),
+        ).run(),
+        expected_msg=(
+            "sqlmesh.utils.errors.TestError: Detected unknown column(s)\n\n"
+            "Expected column(s): id, value\n"
+            "Unknown column(s): foo\n"
         ),
-        test_name="test_foo",
-        model=_create_model("SELECT id, value FROM raw"),
-        context=Context(config=Config(model_defaults=ModelDefaultsConfig(dialect="duckdb"))),
-    ).run()
-
-    expected_error_msg = (
-        "sqlmesh.core.test.definition.TestError: Detected unknown column(s)\n\n"
-        "Expected column(s): id, value\n"
-        "Unknown column(s): foo\n"
     )
-    _check_successful_or_raise(result, expected_msg=expected_error_msg)
 
 
 def test_empty_rows(sushi_context: Context) -> None:
-    model = _create_model(
-        "SELECT id FROM sushi.items", default_catalog=sushi_context.default_catalog
-    )
-    model = t.cast(SqlModel, sushi_context.upsert_model(model))
-
-    body = load_yaml(
-        """
+    _check_successful_or_raise(
+        _create_test(
+            body=load_yaml(
+                """
 test_foo:
   model: sushi.foo
   inputs:
     sushi.items: []
   outputs:
     query: []
-        """
+                """
+            ),
+            test_name="test_foo",
+            model=sushi_context.upsert_model(
+                _create_model(
+                    "SELECT id FROM sushi.items", default_catalog=sushi_context.default_catalog
+                )
+            ),
+            context=sushi_context,
+        ).run()
     )
-    result = _create_test(body, "test_foo", model, sushi_context).run()
-    _check_successful_or_raise(result)
 
 
 @pytest.mark.parametrize("full_model_without_ctes", ["snowflake"], indirect=True)
 def test_normalization(full_model_without_ctes: SqlModel) -> None:
-    body = load_yaml(
-        """
+    normalized_body = _create_test(
+        body=load_yaml(
+            """
 test_foo:
   model: sushi.foo
   inputs:
@@ -704,11 +743,12 @@ test_foo:
   vars:
     start: 2022-01-01
     end: 2022-01-01
-        """
-    )
-
-    context = Context(config=Config(model_defaults=ModelDefaultsConfig(dialect="snowflake")))
-    normalized_body = _create_test(body, "test_foo", full_model_without_ctes, context).body
+            """
+        ),
+        test_name="test_foo",
+        model=full_model_without_ctes,
+        context=Context(config=Config(model_defaults=ModelDefaultsConfig(dialect="snowflake"))),
+    ).body
 
     expected_body = {
         "model": '"MEMORY"."SUSHI"."FOO"',
@@ -795,9 +835,10 @@ def test_test_generation(tmp_path: Path) -> None:
 
 
 def test_source_func() -> None:
-    result = _create_test(
-        body=load_yaml(
-            """
+    _check_successful_or_raise(
+        _create_test(
+            body=load_yaml(
+                """
 test_foo:
   model: xyz
   outputs:
@@ -805,57 +846,52 @@ test_foo:
       - month: 2023-01-01
       - month: 2023-02-01
       - month: 2023-03-01
-            """
-        ),
-        test_name="test_foo",
-        model=_create_model(
-            """
-            SELECT range::DATE AS month
-            FROM RANGE(DATE '2023-01-01', DATE '2023-04-01', INTERVAL 1 MONTH) AS r
-            """
-        ),
-        context=Context(config=Config(model_defaults=ModelDefaultsConfig(dialect="duckdb"))),
-    ).run()
-
-    _check_successful_or_raise(result)
-
-
-def test_nested_data_types() -> None:
-    raw = _create_model(
-        "SELECT array::INT[], struct::STRUCT(x INT[], y VARCHAR, z INT, w STRUCT(a INT)) FROM sushi.unknown",
-        meta="MODEL (name sushi.raw, kind FULL)",
-        default_catalog="memory",
+                """
+            ),
+            test_name="test_foo",
+            model=_create_model(
+                """
+                SELECT range::DATE AS month
+                FROM RANGE(DATE '2023-01-01', DATE '2023-04-01', INTERVAL 1 MONTH) AS r
+                """
+            ),
+            context=Context(config=Config(model_defaults=ModelDefaultsConfig(dialect="duckdb"))),
+        ).run()
     )
-    context = Context(config=Config(model_defaults=ModelDefaultsConfig(dialect="duckdb")))
-    context.upsert_model(raw)
 
-    result = _create_test(
-        body=load_yaml(
-            """
+
+def test_nested_data_types(sushi_context: Context) -> None:
+    _check_successful_or_raise(
+        _create_test(
+            body=load_yaml(
+                """
 test_foo:
   model: sushi.foo
   inputs:
     sushi.raw:
-      - array: [1, 2, 3]
-        struct: {'x': [1, 2, 3], 'y': 'foo', 'z': 1, 'w': {'a': 5}}
-      - array:
-        - 2
-        - 3
-      - array: [0, 4, 1]
+      columns:
+        array: "INT[]"
+        struct: "STRUCT(x INT[], y VARCHAR, z INT, w STRUCT(a INT))"
+      rows:
+        - array: [1, 2, 3]
+          struct: {'x': [1, 2, 3], 'y': 'foo', 'z': 1, 'w': {'a': 5}}
+        - array:
+          - 2
+          - 3
+        - array: [0, 4, 1]
   outputs:
     query:
       - array: [0, 4, 1]
       - array: [1, 2, 3]
         struct: {'x': [1, 2, 3], 'y': 'foo', 'z': 1, 'w': {'a': 5}}
       - array: [2, 3]
-            """
-        ),
-        test_name="test_foo",
-        model=_create_model("SELECT array, struct FROM sushi.raw", default_catalog="memory"),
-        context=context,
-    ).run()
-
-    _check_successful_or_raise(result)
+                """
+            ),
+            test_name="test_foo",
+            model=_create_model("SELECT array, struct FROM sushi.raw", default_catalog="memory"),
+            context=Context(config=Config(model_defaults=ModelDefaultsConfig(dialect="duckdb"))),
+        ).run()
+    )
 
 
 def test_freeze_time(mocker: MockerFixture) -> None:
@@ -899,9 +935,10 @@ test_foo:
 
         return pd.DataFrame([{"ts1": datetime_now, "ts2": current_timestamp}])
 
-    test = _create_test(
-        body=load_yaml(
-            """
+    _check_successful_or_raise(
+        _create_test(
+            body=load_yaml(
+                """
 test_py_model:
   model: py_model
   outputs:
@@ -910,15 +947,14 @@ test_py_model:
         ts2: "2023-01-01 10:05:03"
   vars:
     execution_time: "2023-01-01 12:05:03+02:00"
-            """
-        ),
-        test_name="test_py_model",
-        model=model.get_registry()["py_model"].model(module_path=Path("."), path=Path(".")),
-        context=Context(config=Config(model_defaults=ModelDefaultsConfig(dialect="duckdb"))),
+                """
+            ),
+            test_name="test_py_model",
+            model=model.get_registry()["py_model"].model(module_path=Path("."), path=Path(".")),
+            context=Context(config=Config(model_defaults=ModelDefaultsConfig(dialect="duckdb"))),
+        ).run()
     )
 
-    _check_successful_or_raise(test.run())
-
 
 def test_successes(sushi_context: Context) -> None:
     results = sushi_context.test()