Feat: allow omitting columns in unit tests (#1714)

* Feat: allow omitting columns in unit tests

* Use None instead of np.nan

* Ignore mypy

Author: georgesittas

sqlmesh/core/test/definition.py

@@ -77,7 +77,11 @@ def setUp(self) -> None:
                     v = v.real if hasattr(v, "real") else v
                     columns_to_types[i] = parse_one(type(v).__name__, into=exp.DataType)
 
-            columns_to_types = {k: v for k, v in columns_to_types.items() if k in df}
+            for index, (column, dtype) in enumerate(columns_to_types.items()):
+                if column not in df:
+                    # Fill in missing columns with NULL values
+                    df.insert(index, column, None)  # type: ignore
+
             table = exp.to_table(table_name)
 
             if table.db:

tests/core/test_test.py

@@ -277,6 +277,42 @@ def test_nan(sushi_context: Context, full_model_without_ctes: SqlModel) -> None:
     assert result and result.wasSuccessful()
 
 
+def test_partial_inputs(sushi_context: Context) -> None:
+    model = t.cast(
+        SqlModel,
+        sushi_context.upsert_model(
+            load_sql_based_model(
+                parse(
+                    """
+                    MODEL (
+                        name sushi.foo,
+                        kind FULL
+                    );
+
+                    SELECT id, name FROM sushi.waiter_names;
+                    """,
+                ),
+            ),
+        ),
+    )
+
+    body = load_yaml(
+        """
+test_foo:
+  model: sushi.foo
+  inputs:
+    sushi.waiter_names:
+      - id: 1
+  outputs:
+    query:
+      - id: 1
+        name: null
+            """
+    )
+    result = _create_test(body, "test_foo", model, sushi_context).run()
+    assert result and result.wasSuccessful()
+
+
 @pytest.mark.parametrize("full_model_without_ctes", ["snowflake"], indirect=True)
 def test_normalization(full_model_without_ctes: SqlModel) -> None:
     body = load_yaml(