Feat: add support for CSV input data loading from file or inline (#2640)

Author: themisvaltinos

sqlmesh/core/test/definition.py

@@ -10,6 +10,7 @@
 
 import numpy as np
 import pandas as pd
+from io import StringIO
 from freezegun import freeze_time
 from pandas.api.types import is_object_dtype
 from sqlglot import Dialect, exp
@@ -24,6 +25,7 @@
 from sqlmesh.utils import UniqueKeyDict, random_id, type_is_known, yaml
 from sqlmesh.utils.date import pandas_timestamp_to_pydatetime
 from sqlmesh.utils.errors import ConfigError, TestError
+from sqlmesh.utils.yaml import load as yaml_load
 
 if t.TYPE_CHECKING:
     from sqlglot.dialects.dialect import DialectType
@@ -328,6 +330,17 @@ def _normalize_rows(
             rows = values.get("rows")
             query = values.get("query")
 
+            format = values.get("format")
+            path = values.get("path")
+            if format == "csv":
+                rows = pd.read_csv(path or StringIO(rows)).to_dict(orient="records")
+            elif format in (None, "yaml"):
+                if path:
+                    input_rows = yaml_load(Path(path))
+                    rows = input_rows.get("rows") if isinstance(input_rows, dict) else input_rows
+            else:
+                _raise_error(f"Unsupported data format '{format}' for '{name}'", self.path)
+
             if query is not None:
                 if rows is not None:
                     _raise_error(

tests/core/test_test.py

@@ -345,20 +345,34 @@ def test_row_order(sushi_context: Context, full_model_without_ctes: SqlModel) ->
     )
 
 
-def test_partial_data(sushi_context: Context) -> None:
+@pytest.mark.parametrize(
+    "waiter_names_input",
+    [
+        """sushi.waiter_names:
+          - id: 1
+          - id: 2
+            name: null
+          - id: 3
+            name: 'bob'
+        """,
+        """sushi.waiter_names: 
+      format: csv
+      rows: |
+        id,name
+        1,
+        2,null
+        3,bob""",
+    ],
+)
+def test_partial_data(sushi_context: Context, waiter_names_input: str) -> None:
     _check_successful_or_raise(
         _create_test(
             body=load_yaml(
-                """
+                f"""
 test_foo:
   model: sushi.foo
   inputs:
-    sushi.waiter_names:
-      - id: 1
-      - id: 2
-        name: null
-      - id: 3
-        name: 'bob'
+    {waiter_names_input}
   outputs:
     ctes:
       source:
@@ -390,6 +404,195 @@ def test_partial_data(sushi_context: Context) -> None:
     )
 
 
+@pytest.mark.parametrize(
+    "waiter_names_input",
+    [
+        """sushi.waiter_names:
+        format: yaml
+        rows: 
+        - id: 1
+          name: alice
+        - id: 2
+          name: 'bob'
+        """,
+        """sushi.waiter_names: 
+      format: csv
+      rows: |
+        id,name
+        1,alice
+        2,bob""",
+    ],
+)
+def test_format_inline(sushi_context: Context, waiter_names_input: str) -> None:
+    _check_successful_or_raise(
+        _create_test(
+            body=load_yaml(
+                f"""
+test_foo:
+  model: sushi.foo
+  inputs:
+    {waiter_names_input}
+  outputs:
+    query:
+      - id: 1
+        name: alice
+      - id: 2
+        name: 'bob'
+                """
+            ),
+            test_name="test_foo",
+            model=sushi_context.upsert_model(
+                _create_model(
+                    "SELECT id, name FROM sushi.waiter_names ",
+                    default_catalog=sushi_context.default_catalog,
+                )
+            ),
+            context=sushi_context,
+        ).run()
+    )
+
+
+@pytest.mark.parametrize(
+    ["input_data", "filename", "file_data"],
+    [
+        [
+            """sushi.waiter_names:
+        format: yaml
+        path: """,
+            "test_data.yaml",
+            """- id: 1
+  name: alice
+- id: 2
+  name: 'bob'
+""",
+        ],
+        [
+            """sushi.waiter_names:
+        path: """,
+            "test_data.yaml",
+            """rows:
+- id: 1
+  name: alice
+- id: 2
+  name: 'bob'
+""",
+        ],
+        [
+            """sushi.waiter_names:
+        format: csv
+        path: """,
+            "test_data.csv",
+            """id,name
+1,alice
+2,bob""",
+        ],
+    ],
+)
+def test_format_path(
+    sushi_context: Context, tmp_path: Path, input_data: str, filename: str, file_data: str
+) -> None:
+    test_csv_file = tmp_path / filename
+    test_csv_file.write_text(file_data)
+
+    _check_successful_or_raise(
+        _create_test(
+            body=load_yaml(
+                f"""
+test_foo:
+  model: sushi.foo
+  inputs:
+    {input_data}{str(test_csv_file)}
+  outputs:
+    query:
+      - id: 1
+        name: alice
+      - id: 2
+        name: 'bob'
+                """
+            ),
+            test_name="test_foo",
+            model=sushi_context.upsert_model(
+                _create_model(
+                    "SELECT id, name FROM sushi.waiter_names ",
+                    default_catalog=sushi_context.default_catalog,
+                )
+            ),
+            context=sushi_context,
+        ).run()
+    )
+
+
+def test_unsupported_format_failure(
+    sushi_context: Context, full_model_without_ctes: SqlModel
+) -> None:
+    with pytest.raises(
+        TestError,
+        match="Unsupported data format 'xml' for 'sushi.waiter_names'",
+    ):
+        _create_test(
+            body=load_yaml(
+                """
+test_foo:
+  model: sushi.foo
+  description: XML format isn't supported to load data (fails intentionally)
+  inputs:
+    sushi.waiter_names:
+      format: xml
+      path: 'test_data.xml'
+  outputs:
+    query:
+      - id: 1
+        value: null
+                """
+            ),
+            test_name="test_foo",
+            model=sushi_context.upsert_model(full_model_without_ctes),
+            context=sushi_context,
+        )
+
+    with pytest.raises(
+        TestError,
+        match="Unsupported data format 'xml' for 'sushi.waiter_names'",
+    ):
+        _create_test(
+            body=load_yaml(
+                """
+test_foo:
+  model: sushi.foo
+  description: XML without path doesn't raise error
+  inputs:
+    sushi.waiter_names:
+      format: xml
+      rows: |
+        <rows>
+          <row>
+              <id>1</id>
+              <name>alice</name>
+          </row>
+          <row>
+            <id>2</id>
+            <name>bob</name>
+          </row>
+        </rows>
+  outputs:
+    query:
+      - id: 1
+        name: alice
+      - id: 2
+        name: 'bob'
+                """
+            ),
+            test_name="test_foo",
+            model=sushi_context.upsert_model(
+                _create_model(
+                    "SELECT id, name FROM sushi.waiter_names ",
+                    default_catalog=sushi_context.default_catalog,
+                )
+            ),
+            context=sushi_context,
+        )
+
+
 def test_partial_output_columns() -> None:
     _check_successful_or_raise(
         _create_test(