Make rows key optional in test YAML (#937)

vchan · web-flow · commit 4558b005fee2 · 2023-06-01T15:06:43.000-07:00
diff --git a/sqlmesh/core/test/definition.py b/sqlmesh/core/test/definition.py
@@ -12,6 +12,8 @@
 from sqlmesh.core.model import Model, PythonModel, SqlModel
 from sqlmesh.utils.errors import SQLMeshError
 
+Row = t.Dict[str, t.Any]
+
 
 class TestError(SQLMeshError):
     """Test error"""
@@ -55,7 +57,9 @@ def __init__(
 
     def setUp(self) -> None:
         """Load all input tables"""
-        inputs = {name: table["rows"] for name, table in self.body.get("inputs", {}).items()}
+        inputs = {
+            name: self._get_rows(table) for name, table in self.body.get("inputs", {}).items()
+        }
 
         for table, rows in inputs.items():
             df = pd.DataFrame.from_records(rows)  # noqa
@@ -128,6 +132,18 @@ def create_test(
     def __str__(self) -> str:
         return f"{self.test_name} ({self.path}:{self.body.lc.line})"  # type: ignore
 
+    def _get_rows(self, table: list[Row] | dict[str, list[Row]]) -> list[Row]:
+        """Get a list of rows for input and output table data.
+
+        Input and output table data might be a list of rows or it might be a dictionary
+        with a "rows" key.
+        """
+        if isinstance(table, dict):
+            if "rows" not in table:
+                _raise_error("Incomplete test, missing row data for table", self.path)
+            return table["rows"]
+        return table
+
 
 class SqlModelTest(ModelTest):
     def __init__(
@@ -180,16 +196,16 @@ def test_ctes(self) -> None:
                 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"])
+                expected_df = pd.DataFrame.from_records(self._get_rows(value))
                 actual_df = self.execute(cte_query)
                 self.assert_equal(expected_df, actual_df)
 
     def runTest(self) -> None:
         self.test_ctes()
 
         # Test model query
-        if "rows" in self.body["outputs"].get("query", {}):
-            expected_df = pd.DataFrame.from_records(self.body["outputs"]["query"]["rows"])
+        if "query" in self.body["outputs"]:
+            expected_df = pd.DataFrame.from_records(self._get_rows(self.body["outputs"]["query"]))
             actual_df = self.execute(self.query)
             self.assert_equal(expected_df, actual_df)
 
@@ -224,8 +240,8 @@ def __init__(
         )
 
     def runTest(self) -> None:
-        if "rows" in self.body["outputs"].get("query", {}):
-            expected_df = pd.DataFrame.from_records(self.body["outputs"]["query"]["rows"])
+        if "query" in self.body["outputs"]:
+            expected_df = pd.DataFrame.from_records(self._get_rows(self.body["outputs"]["query"]))
             actual_df = next(
                 self.model.render(
                     context=self.context,
diff --git a/tests/core/test_test.py b/tests/core/test_test.py
@@ -40,18 +40,121 @@ def test_ctes(sushi_context: Context) -> None:
   model: sushi.foo
   inputs:
     raw:
-      rows:
+      - id: 1
+  outputs:
+    ctes:
+      source:
         - id: 1
+      renamed:
+        - fid: 1
+    query:
+      - 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()
+
+
+def test_ctes_only(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:
+      - id: 1
   outputs:
     ctes:
       source:
-        rows:
         - id: 1
       renamed:
-        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()
+
+
+def test_query_only(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:
+      - id: 1
+  outputs:
     query:
-      rows:
       - fid: 1
   vars:
     start: 2022-01-01
@@ -68,3 +171,110 @@ def test_ctes(sushi_context: Context) -> None:
         path=None,
     ).run()
     assert result and result.wasSuccessful()
+
+
+def test_with_rows(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
+        )
+        SELECT id FROM source;
+        """
+                )
+            )
+        ),
+    )
+
+    body = load_yaml(
+        """
+test_foo:
+  model: sushi.foo
+  inputs:
+    raw:
+      rows:
+        - id: 1
+  outputs:
+    ctes:
+      source:
+        rows:
+          - id: 1
+    query:
+      rows:
+        - id: 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()
+
+
+def test_without_rows(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
+        )
+        SELECT id FROM source;
+        """
+                )
+            )
+        ),
+    )
+
+    body = load_yaml(
+        """
+test_foo:
+  model: sushi.foo
+  inputs:
+    raw:
+      - id: 1
+  outputs:
+    ctes:
+      source:
+        - id: 1
+    query:
+      - id: 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()
