Repository: TobikoData/sqlmesh
Pull request #2213 (2024-03-04 20:04:50)

FIx: Refresh files to avoid duplicates


Patch:
diff --git a/sqlmesh/cli/main.py b/sqlmesh/cli/main.py
index dcb813e1..2043e3b6 100644
--- a/sqlmesh/cli/main.py
+++ b/sqlmesh/cli/main.py
@@ -11,7 +11,7 @@ from sqlmesh import configure_logging
 from sqlmesh.cli import error_handler
 from sqlmesh.cli import options as opt
 from sqlmesh.cli.example_project import ProjectTemplate, init_example_project
-from sqlmesh.core.config import Config, load_configs
+from sqlmesh.core.config import load_configs
 from sqlmesh.core.context import Context
 from sqlmesh.utils.date import TimeLike
 from sqlmesh.utils.errors import MissingDependencyError
@@ -77,7 +77,7 @@ def cli(
         elif ctx.invoked_subcommand in ("create_external_models", "migrate", "rollback", "ui"):
             load = False
 
-    configs = load_configs(config, Config, paths)
+    configs = load_configs(config, Context.CONFIG_TYPE, paths)
     log_limit = list(configs.values())[0].log_limit
     configure_logging(debug, ignore_warnings, log_to_stdout, log_limit=log_limit)
 
diff --git a/sqlmesh/core/config/connection.py b/sqlmesh/core/config/connection.py
index 861c4e29..5819fef0 100644
--- a/sqlmesh/core/config/connection.py
+++ b/sqlmesh/core/config/connection.py
@@ -24,6 +24,7 @@ from sqlmesh.core.config.common import (
 from sqlmesh.core.engine_adapter import EngineAdapter
 from sqlmesh.utils.errors import ConfigError
 from sqlmesh.utils.pydantic import (
+    PYDANTIC_MAJOR_VERSION,
     field_validator,
     model_validator,
     model_validator_v1_args,
@@ -42,9 +43,6 @@ class ConnectionConfig(abc.ABC, BaseConfig):
     concurrent_tasks: int
     register_comments: bool
 
-    # Pydantic V2 does not include type field in dict unless also in base class
-    type_: str = Field(alias="type", default="none")
-
     @property
     @abc.abstractmethod
     def _connection_kwargs_keys(self) -> t.Set[str]:
@@ -1276,3 +1274,16 @@ connection_config_validator = field_validator(
     mode="before",
     check_fields=False,
 )(_connection_config_validator)
+
+
+if t.TYPE_CHECKING:
+    # TypeAlias hasn't been introduced until Python 3.10 which means that we can't use it
+    # outside the TYPE_CHECKING guard.
+    SerializableConnectionConfig: t.TypeAlias = ConnectionConfig  # type: ignore
+elif PYDANTIC_MAJOR_VERSION >= 2:
+    import pydantic
+
+    # Workaround for https://docs.pydantic.dev/latest/concepts/serialization/#serializing-with-duck-typing
+    SerializableConnectionConfig = pydantic.SerializeAsAny[ConnectionConfig]  # type: ignore
+else:
+    SerializableConnectionConfig = ConnectionConfig  # type: ignore
diff --git a/sqlmesh/core/config/gateway.py b/sqlmesh/core/config/gateway.py
index 11ab9507..e2ffaff2 100644
--- a/sqlmesh/core/config/gateway.py
+++ b/sqlmesh/core/config/gateway.py
@@ -4,7 +4,10 @@ import typing as t
 
 from sqlmesh.core import constants as c
 from sqlmesh.core.config.base import BaseConfig
-from sqlmesh.core.config.connection import ConnectionConfig, connection_config_validator
+from sqlmesh.core.config.connection import (
+    SerializableConnectionConfig,
+    connection_config_validator,
+)
 from sqlmesh.core.config.scheduler import SchedulerConfig
 
 
@@ -22,9 +25,9 @@ class GatewayConfig(BaseConfig):
             then no schema name is used and therefore the default schema defined for the connection will be used
     """
 
-    connection: t.Optional[ConnectionConfig] = None
-    state_connection: t.Optional[ConnectionConfig] = None
-    test_connection: t.Optional[ConnectionConfig] = None
+    connection: t.Optional[SerializableConnectionConfig] = None
+    state_connection: t.Optional[SerializableConnectionConfig] = None
+    test_connection: t.Optional[SerializableConnectionConfig] = None
     scheduler: t.Optional[SchedulerConfig] = None
     state_schema: t.Optional[str] = c.SQLMESH
 
diff --git a/sqlmesh/core/config/root.py b/sqlmesh/core/config/root.py
index 922877bf..98cc7e71 100644
--- a/sqlmesh/core/config/root.py
+++ b/sqlmesh/core/config/root.py
@@ -17,6 +17,7 @@ from sqlmesh.core.config.base import BaseConfig, UpdateStrategy
 from sqlmesh.core.config.connection import (
     ConnectionConfig,
     DuckDBConnectionConfig,
+    SerializableConnectionConfig,
     connection_config_validator,
 )
 from sqlmesh.core.config.feature_flag import FeatureFlag
@@ -71,8 +72,8 @@ class Config(BaseConfig):
     """
 
     gateways: t.Dict[str, GatewayConfig] = {"": GatewayConfig()}
-    default_connection: ConnectionConfig = DuckDBConnectionConfig()
-    default_test_connection_: t.Optional[ConnectionConfig] = Field(
+    default_connection: SerializableConnectionConfig = DuckDBConnectionConfig()
+    default_test_connection_: t.Optional[SerializableConnectionConfig] = Field(
         default=None, alias="default_test_connection"
     )
     default_scheduler: SchedulerConfig = BuiltInSchedulerConfig()
diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
index 3562b24e..c840b9aa 100644
--- a/sqlmesh/core/context.py
+++ b/sqlmesh/core/context.py
@@ -254,9 +254,10 @@ class GenericContext(BaseContext, t.Generic[C]):
         config_type: The type of config object to use (default Config).
     """
 
+    CONFIG_TYPE: t.Type[C]
+
     def __init__(
         self,
-        config_type: t.Type[C],
         engine_adapter: t.Optional[EngineAdapter] = None,
         notification_targets: t.Optional[t.List[NotificationTarget]] = None,
         state_sync: t.Optional[StateSync] = None,
@@ -271,7 +272,7 @@ class GenericContext(BaseContext, t.Generic[C]):
     ):
         self.console = console or get_console()
         self.configs = (
-            config if isinstance(config, dict) else load_configs(config, config_type, paths)
+            config if isinstance(config, dict) else load_configs(config, self.CONFIG_TYPE, paths)
         )
         self.dag: DAG[str] = DAG()
         self._models: UniqueKeyDict[str, Model] = UniqueKeyDict("models")
@@ -1762,5 +1763,4 @@ class GenericContext(BaseContext, t.Generic[C]):
 
 
 class Context(GenericContext[Config]):
-    def __init__(self, *args: t.Any, **kwargs: t.Any):
-        super().__init__(Config, *args, **kwargs)
+    CONFIG_TYPE = Config
diff --git a/sqlmesh/magics.py b/sqlmesh/magics.py
index f3539db4..74ae9621 100644
--- a/sqlmesh/magics.py
+++ b/sqlmesh/magics.py
@@ -19,7 +19,7 @@ from rich.jupyter import JupyterRenderable
 
 from sqlmesh.cli.example_project import ProjectTemplate, init_example_project
 from sqlmesh.core import constants as c
-from sqlmesh.core.config import Config, load_configs
+from sqlmesh.core.config import load_configs
 from sqlmesh.core.console import get_console
 from sqlmesh.core.context import Context
 from sqlmesh.core.dialect import format_model_expressions, parse
@@ -97,7 +97,7 @@ class SQLMeshMagics(Magics):
         from sqlmesh import configure_logging
 
         args = parse_argstring(self.context, line)
-        configs = load_configs(args.config, Config, args.paths)
+        configs = load_configs(args.config, Context.CONFIG_TYPE, args.paths)
         log_limit = list(configs.values())[0].log_limit
         configure_logging(args.debug, args.ignore_warnings, log_limit=log_limit)
         try:
diff --git a/web/client/src/context/project.ts b/web/client/src/context/project.ts
index 0d2ceda0..d2c1a173 100644
--- a/web/client/src/context/project.ts
+++ b/web/client/src/context/project.ts
@@ -92,9 +92,7 @@ export const useStoreProject = create<ProjectStore>((set, get) => ({
   refreshFiles() {
     const s = get()
 
-    set(() => ({
-      files: new Map(s.files),
-    }))
+    s.setFiles(s.project.allFiles)
   },
 }))
 
diff --git a/web/client/src/library/pages/root/Root.tsx b/web/client/src/library/pages/root/Root.tsx
index 62498f4a..a8313fc9 100644
--- a/web/client/src/library/pages/root/Root.tsx
+++ b/web/client/src/library/pages/root/Root.tsx
@@ -184,6 +184,125 @@ export default function Root(): JSX.Element {
     [synchronizeEnvironment, environment],
   )
 
+  const updateFiles = useCallback(
+    function updateFiles({
+      changes,
+      directories,
+    }: {
+      changes: Array<{
+        change: FileExplorerChange
+        path: string
+        file: ModelFile
+      }>
+      directories: Record<string, Directory>
+    }): void {
+      changes.sort((a: any) =>
+        a.change === EnumFileExplorerChange.Deleted ? -1 : 1,
+      )
+
+      changes.forEach(({ change, path, file }) => {
+        if (change === EnumFileExplorerChange.Modified) {
+          const currentFile = files.get(path)
+
+          if (isNil(currentFile) || isNil(file)) return
+
+          currentFile.update(file)
+        }
+
+        if (change === EnumFileExplorerChange.Deleted) {
+          const artifact = findArtifactByPath(path) ?? files.get(path)
+
+          if (isNil(artifact)) return
+
+          if (artifact instanceof ModelDirectory) {
+            artifact.parent?.removeDirectory(artifact)
+          }
+
+          if (artifact instanceof ModelFile) {
+            artifact.parent?.removeFile(artifact)
+
+            if (inTabs(artifact)) {
+              closeTab(artifact)
+            }
+          }
+        }
+      })
+
+      for (const path in directories) {
+        const directory = directories[path]!
+
+        const currentDirectory = findArtifactByPath(
+          path,
+        ) as Optional<ModelDirectory>
+
+        if (isNil(currentDirectory)) continue
+
+        directory.directories?.forEach(d => {
+          const directory = findArtifactByPath(
+            d.path,
+          ) as Optional<ModelDirectory>
+
+          if (isNil(directory)) {
+            currentDirectory.addDirectory(
+              new ModelDirectory(d, currentDirectory),
+            )
+          }
+        })
+
+        directory.files?.forEach(f => {
+          const file = files.get(f.path)
+
+          if (isNil(file)) {
+            currentDirectory.addFile(new ModelFile(f, currentDirectory))
+          }
+        })
+
+        currentDirectory.directories.sort((a, b) => (a.name > b.name ? 1 : -1))
+        currentDirectory.files.sort((a, b) => (a.name > b.name ? 1 : -1))
+      }
+
+      refreshFiles()
+      setActiveRange()
+    },
+    [files],
+  )
+
+  const formatFile = useCallback(
+    function formatFile(formatFileStatus: FormatFileStatus): void {
+      const file = files.get(formatFileStatus.path)
+
+      if (isNotNil(file)) {
+        file.isFormatted = formatFileStatus.status === Status.success
+
+        refreshFiles()
+      }
+    },
+    [files],
+  )
+
+  const restoreEditorTabsFromSaved = useCallback(
+    function restoreEditorTabsFromSaved(files: ModelFile[]): void {
+      if (isArrayEmpty(storedTabs)) return
+
+      const tabs = storedTabs.map(({ id, content }) => {
+        const file = files.find(file => file.id === id) ?? createLocalFile(id)
+        const storedTab = createTab(file)
+
+        storedTab.file.content = content ?? storedTab.file.content ?? ''
+
+        return storedTab
+      })
+      const tab = tabs.find(tab => tab.file.id === storedTabId)
+
+      addTabs(tabs)
+
+      if (isNotNil(tab) && isNil(selectedFile)) {
+        selectTab(tab)
+      }
+    },
+    [selectTab],
+  )
+
   useEffect(() => {
     void getMeta().then(({ data }) => {
       setVersion(data?.version)
@@ -200,12 +319,8 @@ export default function Root(): JSX.Element {
     const channelErrors = channel('errors', displayErrors)
     const channelTests = channel('tests', updateTests)
     const channelModels = channel('models', updateModels)
-    const channelFile = channel('file', updateFiles)
-    const channelFormatFile = channel('format-file', formatFile)
 
     channelModels.subscribe()
-    channelFormatFile.subscribe()
-    channelFile.subscribe()
 
     if (modules.hasPlans) {
       channelTests.subscribe()
@@ -234,8 +349,6 @@ export default function Root(): JSX.Element {
       void cancelRequestModels()
 
       channelModels.unsubscribe()
-      channelFormatFile?.unsubscribe()
-      channelFile?.unsubscribe()
 
       if (modules.hasPlans) {
         cancelRequestEnvironments()
@@ -304,6 +417,26 @@ export default function Root(): JSX.Element {
     }
   }, [updatePlanCancelTracker])
 
+  useEffect(() => {
+    const channelFile = channel('file', updateFiles)
+
+    channelFile.subscribe()
+
+    return () => {
+      channelFile?.unsubscribe()
+    }
+  }, [updateFiles])
+
+  useEffect(() => {
+    const channelFormatFile = channel('format-file', formatFile)
+
+    channelFormatFile.subscribe()
+
+    return () => {
+      channelFormatFile?.unsubscribe()
+    }
+  }, [formatFile])
+
   useEffect(() => {
     if (location.pathname === EnumRoutes.Home) {
       navigate(modules.defaultNavigationRoute(), { replace: true })
@@ -329,84 +462,6 @@ export default function Root(): JSX.Element {
     setTests(tests)
   }
 
-  function updateFiles({
-    changes,
-    directories,
-  }: {
-    changes: Array<{
-      change: FileExplorerChange
-      path: string
-      file: ModelFile
-    }>
-    directories: Record<string, Directory>
-  }): void {
-    changes.sort((a: any) =>
-      a.change === EnumFileExplorerChange.Deleted ? -1 : 1,
-    )
-
-    changes.forEach(({ change, path, file }) => {
-      if (change === EnumFileExplorerChange.Modified) {
-        const currentFile = files.get(path)
-
-        if (isNil(currentFile) || isNil(file)) return
-
-        currentFile.update(file)
-      }
-
-      if (change === EnumFileExplorerChange.Deleted) {
-        const artifact = findArtifactByPath(path) ?? files.get(path)
-
-        if (isNil(artifact)) return
-
-        if (artifact instanceof ModelDirectory) {
-          artifact.parent?.removeDirectory(artifact)
-        }
-
-        if (artifact instanceof ModelFile) {
-          artifact.parent?.removeFile(artifact)
-
-          if (inTabs(artifact)) {
-            closeTab(artifact)
-          }
-        }
-      }
-    })
-
-    for (const path in directories) {
-      const directory = directories[path]!
-
-      const currentDirectory = findArtifactByPath(path) as
-        | ModelDirectory
-        | undefined
-
-      if (isNil(currentDirectory)) continue
-
-      directory.directories?.forEach((d: any) => {
-        const directory = findArtifactByPath(d.path) as
-          | ModelDirectory
-          | undefined
-
-        if (isNil(directory)) {
-          currentDirectory.addDirectory(new ModelDirectory(d, currentDirectory))
-        }
-      })
-
-      directory.files?.forEach((f: any) => {
-        const file = files.get(path)
-
-        if (isNil(file)) {
-          currentDirectory.addFile(new ModelFile(f, currentDirectory))
-        }
-      })
-
-      currentDirectory.directories.sort((a, b) => (a.name > b.name ? 1 : -1))
-      currentDirectory.files.sort((a, b) => (a.name > b.name ? 1 : -1))
-    }
-
-    refreshFiles()
-    setActiveRange()
-  }
-
   function updateEnviroments(data: Optional<Environments>): void {
     const { environments, default_target_environment, pinned_environments } =
       data ?? {}
@@ -420,42 +475,12 @@ export default function Root(): JSX.Element {
     }
   }
 
-  function formatFile(formatFileStatus: FormatFileStatus): void {
-    const file = files.get(formatFileStatus.path)
-
-    if (isNotNil(file)) {
-      file.isFormatted = formatFileStatus.status === Status.success
-
-      refreshFiles()
-    }
-  }
-
   function closeModalConfirmation(confirmation?: Confirmation): void {
     confirmation?.cancel?.()
 
     setShowConfirmation(false)
   }
 
-  function restoreEditorTabsFromSaved(files: ModelFile[]): void {
-    if (isArrayEmpty(storedTabs)) return
-
-    const tabs = storedTabs.map(({ id, content }) => {
-      const file = files.find(file => file.id === id) ?? createLocalFile(id)
-      const storedTab = createTab(file)
-
-      storedTab.file.content = content ?? storedTab.file.content ?? ''
-
-      return storedTab
-    })
-    const tab = tabs.find(tab => tab.file.id === storedTabId)
-
-    addTabs(tabs)
-
-    if (isNotNil(tab) && isNil(selectedFile)) {
-      selectTab(tab)
-    }
-  }
-
   const confirmation = confirmations[0]
 
   return (


Test patch:
diff --git a/sqlmesh/core/test/definition.py b/sqlmesh/core/test/definition.py
index 33a95779..003dbdbe 100644
--- a/sqlmesh/core/test/definition.py
+++ b/sqlmesh/core/test/definition.py
@@ -6,7 +6,8 @@ from pathlib import Path
 
 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 @@ class ModelTest(unittest.TestCase):
                 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 @@ class ModelTest(unittest.TestCase):
             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
diff --git a/tests/core/test_config.py b/tests/core/test_config.py
index 60494936..23576db4 100644
--- a/tests/core/test_config.py
+++ b/tests/core/test_config.py
@@ -437,3 +437,27 @@ def test_load_alternative_config_type(yaml_config_path: Path, python_config_path
         },
         model_defaults=ModelDefaultsConfig(dialect=""),
     )
+
+
+def test_connection_config_serialization():
+    config = Config(
+        default_connection=DuckDBConnectionConfig(database="my_db"),
+        default_test_connection=DuckDBConnectionConfig(database="my_test_db"),
+    )
+    serialized = config.dict()
+    assert serialized["default_connection"] == {
+        "concurrent_tasks": 1,
+        "register_comments": True,
+        "type": "duckdb",
+        "extensions": [],
+        "connector_config": {},
+        "database": "my_db",
+    }
+    assert serialized["default_test_connection"] == {
+        "concurrent_tasks": 1,
+        "register_comments": True,
+        "type": "duckdb",
+        "extensions": [],
+        "connector_config": {},
+        "database": "my_test_db",
+    }
diff --git a/tests/core/test_test.py b/tests/core/test_test.py
index 4af267e1..ffa2f5bc 100644
--- a/tests/core/test_test.py
+++ b/tests/core/test_test.py
@@ -16,6 +16,9 @@ from sqlmesh.core.test.definition import SqlModelTest
 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 @@ test_foo:
         """
     )
     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 @@ test_foo:
         """
     )
     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 @@ test_foo:
         """
     )
     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 @@ test_foo:
         """
     )
     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 @@ test_foo:
         """
     )
     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 @@ test_foo:
         """
     )
     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,7 +281,7 @@ test_foo:
 
     # 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()
@@ -274,9 +289,16 @@ test_foo:
     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 @@ test_foo:
         """
     )
     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 @@ test_foo:
         """
     )
     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 @@ test_foo:
         """
     )
     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 @@ test_foo:
         """
     )
     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 @@ test_foo:
         """
     )
     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 @@ test_foo:
 
 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 @@ test_foo:
         """
     )
     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 @@ test_foo:
         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)