Include a model's CTEs when testing one of its CTEs (#920)

Author: vchan

sqlmesh/core/test/definition.py

@@ -18,6 +18,7 @@ class TestError(SQLMeshError):
 
 
 class ModelTest(unittest.TestCase):
+    __test__ = False
     view_names: list[str] = []
 
     def __init__(
@@ -156,7 +157,9 @@ def __init__(
             engine_adapter=engine_adapter,
         )
         # For tests we just use the model name for the table reference and we don't want to expand
-        mapping = {name: _test_fixture_name(name) for name in models}
+        mapping = {
+            name: _test_fixture_name(name) for name in models.keys() | body.get("inputs", {}).keys()
+        }
         if mapping:
             self.query = exp.replace_tables(self.query, mapping)
 
@@ -174,8 +177,11 @@ def test_ctes(self) -> None:
                     _raise_error(
                         f"No CTE named {cte_name} found in model {self.model.name}", self.path
                     )
+                cte_query = self.ctes[cte_name].this
+                for alias, cte in self.ctes.items():
+                    cte_query = cte_query.with_(alias, cte.this)
                 expected_df = pd.DataFrame.from_records(value["rows"])
-                actual_df = self.execute(self.ctes[cte_name].this)
+                actual_df = self.execute(cte_query)
                 self.assert_equal(expected_df, actual_df)
 
     def runTest(self) -> None:

tests/core/test_test.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+import typing as t
+
+from sqlmesh.core.context import Context
+from sqlmesh.core.dialect import parse
+from sqlmesh.core.model import SqlModel, load_model
+from sqlmesh.core.test.definition import SqlModelTest
+from sqlmesh.utils.yaml import load as load_yaml
+
+
+def test_ctes(sushi_context: Context) -> None:
+    model = t.cast(
+        SqlModel,
+        sushi_context.upsert_model(
+            load_model(
+                parse(
+                    """
+        MODEL (
+            name sushi.foo,
+            kind FULL,
+        );
+
+        WITH source AS (
+            SELECT id FROM raw
+        ),
+        renamed AS (
+            SELECT id as fid FROM source
+        )
+        SELECT fid FROM renamed;
+        """
+                )
+            )
+        ),
+    )
+
+    body = load_yaml(
+        """
+test_foo:
+  model: sushi.foo
+  inputs:
+    raw:
+      rows:
+        - id: 1
+  outputs:
+    ctes:
+      source:
+        rows:
+        - id: 1
+      renamed:
+        rows:
+        - fid: 1
+    query:
+      rows:
+      - fid: 1
+  vars:
+    start: 2022-01-01
+    end: 2022-01-01
+    latest: 2022-01-01
+            """
+    )
+    result = SqlModelTest(
+        body=body["test_foo"],
+        test_name="test_foo",
+        model=model,
+        models=sushi_context._models,
+        engine_adapter=sushi_context._test_engine_adapter,
+        path=None,
+    ).run()
+    assert result and result.wasSuccessful()