Feat: add support for arrays in unit tests (#2215)

* Feat: add support for arrays in unit tests

* Fix formatting

* Use apply instead of map

* Fix test

* Test itemized YAML format for array as well

Author: georgesittas

sqlmesh/core/test/definition.py

@@ -6,7 +6,8 @@
 
 import numpy as np
 import pandas as pd
-from sqlglot import exp, parse_one
+from sqlglot import exp
+from sqlglot.optimizer.annotate_types import annotate_types
 from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
 
 from sqlmesh.core import constants as c
@@ -82,7 +83,8 @@ def setUp(self) -> None:
                 for i, v in rows[0].items():
                     # convert ruamel into python
                     v = v.real if hasattr(v, "real") else v
-                    columns_to_types[i] = parse_one(type(v).__name__, into=exp.DataType)
+                    v_type = annotate_types(exp.convert(v)).type or type(v).__name__
+                    columns_to_types[i] = exp.maybe_parse(v_type, into=exp.DataType)
 
             test_fixture_table = _fully_qualified_test_fixture_table(table_name, self.dialect)
             if test_fixture_table.db:
@@ -112,21 +114,28 @@ def assert_equal(self, expected: pd.DataFrame, actual: pd.DataFrame, sort: bool)
             actual_types, errors="ignore"
         )
 
-        expected = expected.replace({None: np.nan})
         actual = actual.replace({None: np.nan})
+        expected = expected.replace({None: np.nan})
+
+        def _to_hashable(x: t.Any) -> t.Any:
+            return tuple(x) if isinstance(x, list) else x
 
         try:
+            if sort:
+                actual = (
+                    actual.apply(_to_hashable)
+                    .sort_values(by=actual.columns.to_list())
+                    .reset_index(drop=True)
+                )
+                expected = (
+                    expected.apply(_to_hashable)
+                    .sort_values(by=expected.columns.to_list())
+                    .reset_index(drop=True)
+                )
+
             pd.testing.assert_frame_equal(
-                (
-                    expected.sort_values(by=expected.columns.to_list()).reset_index(drop=True)
-                    if sort
-                    else expected
-                ),
-                (
-                    actual.sort_values(by=actual.columns.to_list()).reset_index(drop=True)
-                    if sort
-                    else actual
-                ),
+                expected,
+                actual,
                 check_dtype=False,
                 check_datetimelike_compat=True,
                 check_like=True,  # ignore column order

tests/core/test_test.py

@@ -16,6 +16,9 @@
 from sqlmesh.utils.errors import ConfigError
 from sqlmesh.utils.yaml import load as load_yaml
 
+if t.TYPE_CHECKING:
+    from unittest import TestResult
+
 pytestmark = pytest.mark.slow
 
 SUSHI_FOO_META = "MODEL (name sushi.foo, kind FULL)"
@@ -52,6 +55,18 @@ def _create_model(
     )
 
 
+def _check_successful_or_raise(
+    result: t.Optional[TestResult], expected_failure_msg: t.Optional[str] = None
+) -> None:
+    assert result is not None
+    if not result.wasSuccessful():
+        error_or_failure_traceback = (result.errors or result.failures)[0][1]
+        if result.failures and expected_failure_msg:
+            assert expected_failure_msg in error_or_failure_traceback
+        else:
+            raise AssertionError(error_or_failure_traceback)
+
+
 @pytest.fixture
 def full_model_without_ctes(request) -> SqlModel:
     return _create_model(
@@ -110,7 +125,7 @@ def test_ctes(sushi_context: Context, full_model_with_two_ctes: SqlModel) -> Non
         """
     )
     result = _create_test(body, "test_foo", model, sushi_context).run()
-    assert result and result.wasSuccessful()
+    _check_successful_or_raise(result)
 
 
 def test_ctes_only(sushi_context: Context, full_model_with_two_ctes: SqlModel) -> None:
@@ -134,7 +149,7 @@ def test_ctes_only(sushi_context: Context, full_model_with_two_ctes: SqlModel) -
         """
     )
     result = _create_test(body, "test_foo", model, sushi_context).run()
-    assert result and result.wasSuccessful()
+    _check_successful_or_raise(result)
 
 
 def test_query_only(sushi_context: Context, full_model_with_two_ctes: SqlModel) -> None:
@@ -155,7 +170,7 @@ def test_query_only(sushi_context: Context, full_model_with_two_ctes: SqlModel)
         """
     )
     result = _create_test(body, "test_foo", model, sushi_context).run()
-    assert result and result.wasSuccessful()
+    _check_successful_or_raise(result)
 
 
 def test_with_rows(sushi_context: Context, full_model_with_single_cte: SqlModel) -> None:
@@ -182,7 +197,7 @@ def test_with_rows(sushi_context: Context, full_model_with_single_cte: SqlModel)
         """
     )
     result = _create_test(body, "test_foo", model, sushi_context).run()
-    assert result and result.wasSuccessful()
+    _check_successful_or_raise(result)
 
 
 def test_without_rows(sushi_context: Context, full_model_with_single_cte: SqlModel) -> None:
@@ -206,7 +221,7 @@ def test_without_rows(sushi_context: Context, full_model_with_single_cte: SqlMod
         """
     )
     result = _create_test(body, "test_foo", model, sushi_context).run()
-    assert result and result.wasSuccessful()
+    _check_successful_or_raise(result)
 
 
 def test_column_order(sushi_context: Context, full_model_without_ctes: SqlModel) -> None:
@@ -231,7 +246,7 @@ def test_column_order(sushi_context: Context, full_model_without_ctes: SqlModel)
         """
     )
     result = _create_test(body, "test_foo", model, sushi_context).run()
-    assert result and result.wasSuccessful()
+    _check_successful_or_raise(result)
 
 
 def test_row_order(sushi_context: Context, full_model_without_ctes: SqlModel) -> None:
@@ -266,17 +281,24 @@ def test_row_order(sushi_context: Context, full_model_without_ctes: SqlModel) ->
 
     # model query without ORDER BY should pass unit test
     result = _create_test(body, "test_foo", model, sushi_context).run()
-    assert result and result.wasSuccessful()
+    _check_successful_or_raise(result)
 
     # model query with ORDER BY should fail unit test
     full_model_without_ctes_dict = full_model_without_ctes.dict()
     full_model_without_ctes_dict["query"] = full_model_without_ctes.query.order_by("id")  # type: ignore
     full_model_without_ctes_orderby = SqlModel(**full_model_without_ctes_dict)
 
     model = t.cast(SqlModel, sushi_context.upsert_model(full_model_without_ctes_orderby))
-
     result = _create_test(body, "test_foo", model, sushi_context).run()
-    assert result and not result.wasSuccessful()
+
+    expected_failure_msg = """AssertionError: Data differs (exp: expected, act: actual)
+
+   id     value      ds    
+  exp act   exp act exp act
+0   2   1     3   2   4   3
+1   1   2     2   3   3   4"""
+
+    _check_successful_or_raise(result, expected_failure_msg=expected_failure_msg)
 
 
 def test_partial_data(sushi_context: Context) -> None:
@@ -316,7 +338,7 @@ def test_partial_data(sushi_context: Context) -> None:
         """
     )
     result = _create_test(body, "test_foo", model, sushi_context).run()
-    assert result and result.wasSuccessful()
+    _check_successful_or_raise(result)
 
 
 def test_partial_data_column_order(sushi_context: Context) -> None:
@@ -347,7 +369,7 @@ def test_partial_data_column_order(sushi_context: Context) -> None:
         """
     )
     result = _create_test(body, "test_foo", model, sushi_context).run()
-    assert result and result.wasSuccessful()
+    _check_successful_or_raise(result)
 
 
 def test_partial_data_missing_schemas(sushi_context: Context) -> None:
@@ -371,7 +393,7 @@ def test_partial_data_missing_schemas(sushi_context: Context) -> None:
         """
     )
     result = _create_test(body, "test_foo", model, sushi_context).run()
-    assert result and result.wasSuccessful()
+    _check_successful_or_raise(result)
 
     model = _create_model(
         "SELECT *, DATE_TRUNC('month', date)::DATE AS month, NULL::DATE AS null_date, FROM unknown"
@@ -401,7 +423,7 @@ def test_partial_data_missing_schemas(sushi_context: Context) -> None:
         """
     )
     result = _create_test(body, "test_foo", model, sushi_context).run()
-    assert result and result.wasSuccessful()
+    _check_successful_or_raise(result)
 
 
 def test_missing_column_failure(sushi_context: Context, full_model_without_ctes: SqlModel) -> None:
@@ -423,9 +445,8 @@ def test_missing_column_failure(sushi_context: Context, full_model_without_ctes:
         """
     )
     result = _create_test(body, "test_foo", model, sushi_context).run()
-    assert result and not result.wasSuccessful()
 
-    expected_msg = """AssertionError: Data differs (exp: expected, act: actual)
+    expected_failure_msg = """AssertionError: Data differs (exp: expected, act: actual)
 
   value      ds    
     exp act exp act
@@ -434,7 +455,7 @@ def test_missing_column_failure(sushi_context: Context, full_model_without_ctes:
 
 Test description: sushi.foo's output has a missing column (fails intentionally)
 """
-    assert expected_msg in result.failures[0][1]
+    _check_successful_or_raise(result, expected_failure_msg=expected_failure_msg)
 
 
 def test_empty_rows(sushi_context: Context) -> None:
@@ -454,7 +475,7 @@ def test_empty_rows(sushi_context: Context) -> None:
         """
     )
     result = _create_test(body, "test_foo", model, sushi_context).run()
-    assert result and result.wasSuccessful()
+    _check_successful_or_raise(result)
 
 
 @pytest.mark.parametrize("full_model_without_ctes", ["snowflake"], indirect=True)
@@ -541,7 +562,7 @@ def test_test_generation(tmp_path: Path) -> None:
     assert test["test_full_model"]["vars"] == {"start": "2020-01-01", "end": "2024-01-01"}
 
     result = context.test()
-    assert result and result.wasSuccessful()
+    _check_successful_or_raise(result)
 
     context.create_test(
         "sqlmesh_example.full_model", input_queries=input_queries, name="new_name", path="foo/bar"
@@ -557,13 +578,13 @@ def test_source_func() -> None:
         body=load_yaml(
             """
 test_foo:
-    model: xyz
-    outputs:
-        query:
-            - month: 2023-01-01
-            - month: 2023-02-01
-            - month: 2023-03-01
-"""
+  model: xyz
+  outputs:
+    query:
+      - month: 2023-01-01
+      - month: 2023-02-01
+      - month: 2023-03-01
+            """
         ),
         test_name="test_foo",
         model=_create_model(
@@ -575,4 +596,32 @@ def test_source_func() -> None:
         context=Context(config=Config(model_defaults=ModelDefaultsConfig(dialect="duckdb"))),
     ).run()
 
-    assert result and result.wasSuccessful()
+    _check_successful_or_raise(result)
+
+
+def test_nested_data_types() -> None:
+    result = _create_test(
+        body=load_yaml(
+            """
+test_foo:
+  model: sushi.foo
+  inputs:
+    raw:
+      - value: [1, 2, 3]
+      - value:
+        - 2
+        - 3
+      - value: [0, 4, 1]
+  outputs:
+    query:
+      - value: [0, 4, 1]
+      - value: [1, 2, 3]
+      - value: [2, 3]
+            """
+        ),
+        test_name="test_foo",
+        model=_create_model("SELECT value FROM raw"),
+        context=Context(config=Config(model_defaults=ModelDefaultsConfig(dialect="duckdb"))),
+    ).run()
+
+    _check_successful_or_raise(result)