sqlmesh.core.test.definition

   1from __future__ import annotations
   2
   3import sys
   4
   5import datetime
   6import threading
   7import typing as t
   8import unittest
   9from collections import Counter
  10from contextlib import nullcontext, contextmanager, AbstractContextManager
  11from itertools import chain
  12from pathlib import Path
  13from unittest.mock import patch
  14
  15
  16from io import StringIO
  17from sqlglot import Dialect, exp
  18from sqlglot.optimizer.annotate_types import annotate_types
  19from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
  20
  21from sqlmesh.core import constants as c
  22from sqlmesh.core.dialect import normalize_model_name, schema_
  23from sqlmesh.core.engine_adapter import EngineAdapter
  24from sqlmesh.core.macros import RuntimeStage
  25from sqlmesh.core.model import Model, PythonModel, SqlModel
  26from sqlmesh.utils import UniqueKeyDict, random_id, type_is_known, yaml
  27from sqlmesh.utils.date import date_dict, pandas_timestamp_to_pydatetime, to_datetime
  28from sqlmesh.utils.errors import ConfigError, TestError
  29from sqlmesh.utils.yaml import load as yaml_load
  30from sqlmesh.utils import Verbosity
  31from sqlmesh.utils.rich import df_to_table
  32
  33if t.TYPE_CHECKING:
  34    import pandas as pd
  35
  36    from sqlglot.dialects.dialect import DialectType
  37
  38    Row = t.Dict[str, t.Any]
  39
  40
  41TIME_KWARG_KEYS = {
  42    "start",
  43    "end",
  44    "execution_time",
  45    "latest",
  46    # all built-in datetime macro var names
  47    *date_dict(execution_time="1970-01-01", start="1970-01-01", end="1970-01-01").keys(),
  48}
  49
  50
  51class ModelTest(unittest.TestCase):
  52    __test__ = False
  53
  54    CONCURRENT_RENDER_LOCK = threading.Lock()
  55
  56    def __init__(
  57        self,
  58        body: t.Dict[str, t.Any],
  59        test_name: str,
  60        model: Model,
  61        models: UniqueKeyDict[str, Model],
  62        engine_adapter: EngineAdapter,
  63        dialect: str | None = None,
  64        path: Path | None = None,
  65        preserve_fixtures: bool = False,
  66        default_catalog: str | None = None,
  67        concurrency: bool = False,
  68        verbosity: Verbosity = Verbosity.DEFAULT,
  69    ) -> None:
  70        """ModelTest encapsulates a unit test for a model.
  71
  72        Args:
  73            body: A dictionary that contains test metadata like inputs and outputs.
  74            test_name: The name of the test.
  75            model: The model that is being tested.
  76            models: All models to use for expansion and mapping of physical locations.
  77            engine_adapter: The engine adapter to use.
  78            dialect: The models' dialect, used for normalization purposes.
  79            path: An optional path to the test definition yaml file.
  80            preserve_fixtures: Preserve the fixture tables in the testing database, useful for debugging.
  81        """
  82        self.body = body
  83        self.test_name = test_name
  84        self.model = model
  85        self.models = models
  86        self.engine_adapter = engine_adapter
  87        self.path = path
  88        self.preserve_fixtures = preserve_fixtures
  89        self.default_catalog = default_catalog
  90        self.dialect = dialect
  91        self.concurrency = concurrency
  92        self.verbosity = verbosity
  93
  94        self._fixture_table_cache: t.Dict[str, exp.Table] = {}
  95        self._normalized_column_name_cache: t.Dict[str, str] = {}
  96        self._normalized_model_name_cache: t.Dict[t.Tuple[str, bool], str] = {}
  97
  98        self._test_adapter_dialect = Dialect.get_or_raise(self.engine_adapter.dialect)
  99
 100        self._validate_and_normalize_test()
 101
 102        if self.engine_adapter.default_catalog:
 103            self._fixture_catalog: t.Optional[exp.Identifier] = normalize_identifiers(
 104                exp.parse_identifier(
 105                    self.engine_adapter.default_catalog, dialect=self._test_adapter_dialect
 106                ),
 107                dialect=self._test_adapter_dialect,
 108            )
 109        else:
 110            self._fixture_catalog = None
 111
 112        # The test schema name is randomized to avoid concurrency issues,
 113        # unless a schema is provided in the unit tests's body
 114        self._fixture_schema = exp.parse_identifier(
 115            self.body.get("schema") or f"sqlmesh_test_{random_id(short=True)}"
 116        )
 117        self._qualified_fixture_schema = schema_(self._fixture_schema, self._fixture_catalog)
 118
 119        self._transforms = self._test_adapter_dialect.generator_class.TRANSFORMS
 120        self._execution_time = str(self.body.get("vars", {}).get("execution_time") or "")
 121
 122        if self._execution_time:
 123            # Normalizes the execution time by converting it into UTC timezone
 124            self._execution_time = str(to_datetime(self._execution_time))
 125
 126        # When execution_time is set, we mock the CURRENT_* SQL expressions so they always return it
 127        if self._execution_time:
 128            exec_time = exp.Literal.string(self._execution_time)
 129            self._transforms = {
 130                **self._transforms,
 131                exp.CurrentDate: lambda self, _: self.sql(
 132                    exp.cast(exec_time, "date", dialect=dialect)
 133                ),
 134                exp.CurrentDatetime: lambda self, _: self.sql(
 135                    exp.cast(exec_time, "datetime", dialect=dialect)
 136                ),
 137                exp.CurrentTime: lambda self, _: self.sql(
 138                    exp.cast(exec_time, "time", dialect=dialect)
 139                ),
 140                exp.CurrentTimestamp: lambda self, _: self.sql(
 141                    exp.cast(exec_time, "timestamp", dialect=dialect)
 142                ),
 143            }
 144
 145        super().__init__()
 146
 147    def defaultTestResult(self) -> unittest.TestResult:
 148        from sqlmesh.core.test.result import ModelTextTestResult
 149
 150        return ModelTextTestResult(stream=sys.stdout, descriptions=True, verbosity=self.verbosity)
 151
 152    def shortDescription(self) -> t.Optional[str]:
 153        return self.body.get("description")
 154
 155    def setUp(self) -> None:
 156        """Load all input tables"""
 157        import pandas as pd
 158        import numpy as np
 159
 160        self.engine_adapter.create_schema(self._qualified_fixture_schema)
 161
 162        for name, values in self.body.get("inputs", {}).items():
 163            all_types_are_known = False
 164            columns_to_known_types: t.Dict[str, exp.DataType] = {}
 165
 166            model = self.models.get(name)
 167            if model:
 168                inferred_columns_to_types = model.columns_to_types or {}
 169                columns_to_known_types = {
 170                    c: t for c, t in inferred_columns_to_types.items() if type_is_known(t)
 171                }
 172                all_types_are_known = bool(inferred_columns_to_types) and (
 173                    len(columns_to_known_types) == len(inferred_columns_to_types)
 174                )
 175
 176            # Types specified in the test will override the corresponding inferred ones
 177            columns_to_known_types.update(values.get("columns", {}))
 178
 179            rows = values.get("rows")
 180            if not all_types_are_known and rows:
 181                for col, value in rows[0].items():
 182                    if col not in columns_to_known_types:
 183                        v_type = annotate_types(exp.convert(value)).type or type(value).__name__
 184                        v_type = exp.maybe_parse(
 185                            v_type, into=exp.DataType, dialect=self._test_adapter_dialect
 186                        )
 187
 188                        if not type_is_known(v_type):
 189                            _raise_error(
 190                                f"Failed to infer the data type of column '{col}' for '{name}'. This issue can be "
 191                                "mitigated by casting the column in the model definition, setting its type in "
 192                                "external_models.yaml if it's an external model, setting the model's 'columns' property, "
 193                                "or setting its 'columns' mapping in the test itself",
 194                                self.path,
 195                            )
 196
 197                        columns_to_known_types[col] = v_type
 198
 199            if rows is None:
 200                query_or_df: exp.Query | pd.DataFrame = self._add_missing_columns(
 201                    values["query"], columns_to_known_types
 202                )
 203                if columns_to_known_types:
 204                    columns_to_known_types = {
 205                        col: columns_to_known_types[col] for col in query_or_df.named_selects
 206                    }
 207            else:
 208                query_or_df = self._create_df(values, columns=columns_to_known_types)
 209
 210            # Convert NaN/NaT values to None if DataFrame
 211            if isinstance(query_or_df, pd.DataFrame):
 212                query_or_df = query_or_df.replace({np.nan: None})
 213
 214            self.engine_adapter.create_view(
 215                self._test_fixture_table(name), query_or_df, columns_to_known_types
 216            )
 217
 218    def tearDown(self) -> None:
 219        """Drop all fixture tables."""
 220        if not self.preserve_fixtures:
 221            self.engine_adapter.drop_schema(self._qualified_fixture_schema, cascade=True)
 222
 223    def assert_equal(
 224        self,
 225        expected: pd.DataFrame,
 226        actual: pd.DataFrame,
 227        sort: bool,
 228        partial: t.Optional[bool] = False,
 229    ) -> None:
 230        """Compare two DataFrames"""
 231        import numpy as np
 232        import pandas as pd
 233        from pandas.api.types import is_object_dtype
 234
 235        if partial:
 236            intersection = actual[actual.columns.intersection(expected.columns)]
 237            if len(intersection.columns) > 0:
 238                actual = intersection
 239
 240        # Two astypes are necessary, pandas converts strings to times as NS,
 241        # but if the actual is US, it doesn't take effect until the 2nd try!
 242        actual_types = actual.dtypes.to_dict()
 243        expected = expected.astype(actual_types, errors="ignore").astype(
 244            actual_types, errors="ignore"
 245        )
 246
 247        # The `actual` df's dtypes will almost always be pd.Timestamp for datetime values,
 248        # but in some scenarios (e.g., DuckDB >=0.10.2) it will be a pandas `object` type
 249        # containing python `datetime.xxx` values.
 250        #
 251        # Pandas `object` columns result in a noop for the `astype` call above. Because any
 252        # quoted YAML value is a string, we must manually convert the `expected` df string
 253        # values to the correct `datetime.xxx` type.
 254        #
 255        # We determine the type from a single sentinel value, but since the `actual` df is
 256        # coming from a database query, it is safe to assume that the column contains only
 257        # a single type.
 258        object_sentinel_values = {
 259            col: actual[col][0]
 260            for col in actual_types
 261            if is_object_dtype(actual_types[col]) and len(actual[col]) != 0
 262        }
 263        for col, value in object_sentinel_values.items():
 264            try:
 265                # can't use `isinstance()` here - https://stackoverflow.com/a/68743663/1707525
 266                if type(value) is datetime.date:
 267                    expected[col] = pd.to_datetime(expected[col]).dt.date
 268                elif type(value) is datetime.time:
 269                    expected[col] = pd.to_datetime(expected[col]).dt.time
 270                elif type(value) is datetime.datetime:
 271                    expected[col] = pd.to_datetime(expected[col]).dt.to_pydatetime()
 272            except Exception as e:
 273                from sqlmesh.core.console import get_console
 274
 275                get_console().log_warning(
 276                    f"Failed to convert expected value for {col} into `datetime` "
 277                    f"for unit test '{str(self)}'. {str(e)}."
 278                )
 279
 280        actual = actual.replace({np.nan: None})
 281        expected = expected.replace({np.nan: None})
 282
 283        # We define this here to avoid a top-level import of numpy and pandas
 284        DATETIME_TYPES = (
 285            datetime.datetime,
 286            datetime.date,
 287            datetime.time,
 288            np.datetime64,
 289            pd.Timestamp,
 290        )
 291
 292        def _to_hashable(x: t.Any) -> t.Any:
 293            if isinstance(x, (list, np.ndarray)):
 294                return tuple(_to_hashable(v) for v in x)
 295            if isinstance(x, dict):
 296                return tuple((k, _to_hashable(v)) for k, v in x.items())
 297            return str(x) if isinstance(x, DATETIME_TYPES) or not isinstance(x, t.Hashable) else x
 298
 299        actual = actual.apply(lambda col: col.map(_to_hashable))
 300        expected = expected.apply(lambda col: col.map(_to_hashable))
 301
 302        if sort:
 303            actual = actual.sort_values(by=actual.columns.to_list()).reset_index(drop=True)
 304            expected = expected.sort_values(by=expected.columns.to_list()).reset_index(drop=True)
 305
 306        try:
 307            pd.testing.assert_frame_equal(
 308                expected,
 309                actual,
 310                check_dtype=False,
 311                check_datetimelike_compat=True,
 312                check_like=True,  # Ignore column order
 313            )
 314        except AssertionError as e:
 315            # There are 2 concepts at play here:
 316            # 1. The Exception args will contain the error message plus the diff dataframe table stringified
 317            #    (backwards compatibility with existing tests, possible to serialize/send over network etc)
 318            # 2. Each test will also transform these diff dataframes into Rich tables, which will be the ones that'll
 319            #    be surfaced to the user through Console for better UX (versus stringified dataframes)
 320            #
 321            # This is a bit of a hack, but it's a way to get the best of both worlds.
 322            args: t.List[t.Any] = []
 323
 324            failed_subtest = ""
 325
 326            if subtest := getattr(self, "_subtest", None):
 327                if cte := subtest.params.get("cte"):
 328                    failed_subtest = f" (CTE {cte})"
 329
 330            if expected.shape != actual.shape:
 331                _raise_if_unexpected_columns(expected.columns, actual.columns)
 332
 333                args.append("Data mismatch (rows are different)")
 334
 335                missing_rows = _row_difference(expected, actual)
 336                if not missing_rows.empty:
 337                    args[0] += f"\n\nMissing rows:\n\n{missing_rows}"
 338                    args.append(df_to_table(f"Missing rows{failed_subtest}", missing_rows))
 339
 340                unexpected_rows = _row_difference(actual, expected)
 341
 342                if not unexpected_rows.empty:
 343                    args[0] += f"\n\nUnexpected rows:\n\n{unexpected_rows}"
 344                    args.append(df_to_table(f"Unexpected rows{failed_subtest}", unexpected_rows))
 345
 346            else:
 347                diff = expected.compare(actual).rename(columns={"self": "exp", "other": "act"})
 348
 349                args.append(f"Data mismatch (exp: expected, act: actual)\n\n{diff}")
 350
 351                diff.rename(columns={"exp": "Expected", "act": "Actual"}, inplace=True)
 352                if self.verbosity == Verbosity.DEFAULT:
 353                    args.extend(
 354                        df_to_table(f"Data mismatch{failed_subtest}", df)
 355                        for df in _split_df_by_column_pairs(diff)
 356                    )
 357                else:
 358                    from pandas import DataFrame, MultiIndex
 359
 360                    levels = t.cast(MultiIndex, diff.columns).levels[0]
 361                    for col in levels:
 362                        # diff[col] returns a DataFrame when columns is a MultiIndex
 363                        col_diff = t.cast(DataFrame, diff[col])
 364                        if not col_diff.empty:
 365                            table = df_to_table(
 366                                f"[bold red]Column '{col}' mismatch{failed_subtest}[/bold red]",
 367                                col_diff,
 368                            )
 369                            args.append(table)
 370
 371            e.args = (*args,)
 372
 373            raise e
 374
 375    def runTest(self) -> None:
 376        raise NotImplementedError
 377
 378    def path_relative_to(self, other: Path) -> Path | None:
 379        """Compute a version of this test's path relative to the `other` path"""
 380        return self.path.relative_to(other) if self.path else None
 381
 382    @staticmethod
 383    def create_test(
 384        body: t.Dict[str, t.Any],
 385        test_name: str,
 386        models: UniqueKeyDict[str, Model],
 387        engine_adapter: EngineAdapter,
 388        dialect: str | None,
 389        path: Path | None,
 390        preserve_fixtures: bool = False,
 391        default_catalog: str | None = None,
 392        concurrency: bool = False,
 393        verbosity: Verbosity = Verbosity.DEFAULT,
 394    ) -> t.Optional[ModelTest]:
 395        """Create a SqlModelTest or a PythonModelTest.
 396
 397        Args:
 398            body: A dictionary that contains test metadata like inputs and outputs.
 399            test_name: The name of the test.
 400            models: All models to use for expansion and mapping of physical locations.
 401            engine_adapter: The engine adapter to use.
 402            dialect: The models' dialect, used for normalization purposes.
 403            path: An optional path to the test definition yaml file.
 404            preserve_fixtures: Preserve the fixture tables in the testing database, useful for debugging.
 405        """
 406        name = body.get("model")
 407        if name is None:
 408            _raise_error("Missing required 'model' field", path)
 409
 410        name = normalize_model_name(name, default_catalog=default_catalog, dialect=dialect)
 411        model = models.get(name)
 412        if not model:
 413            from sqlmesh.core.console import get_console
 414
 415            get_console().log_warning(
 416                f"Model '{name}' was not found{' at ' + str(path) if path else ''}"
 417            )
 418            return None
 419
 420        if isinstance(model, SqlModel):
 421            test_type: t.Type[ModelTest] = SqlModelTest
 422        elif isinstance(model, PythonModel):
 423            test_type = PythonModelTest
 424        else:
 425            _raise_error(f"Model '{name}' is an unsupported model type for testing", path)
 426
 427        try:
 428            return test_type(
 429                body,
 430                test_name,
 431                t.cast(Model, model),
 432                models,
 433                engine_adapter,
 434                dialect,
 435                path,
 436                preserve_fixtures,
 437                default_catalog,
 438                concurrency,
 439                verbosity,
 440            )
 441        except Exception as e:
 442            raise TestError(f"Failed to create test {test_name} ({path})\n{str(e)}")
 443
 444    def __str__(self) -> str:
 445        return f"{self.test_name} ({self.path})"
 446
 447    def _validate_and_normalize_test(self) -> None:
 448        inputs = self.body.get("inputs")
 449        outputs = self.body.get("outputs", {})
 450
 451        if not outputs:
 452            _raise_error("Incomplete test, missing outputs", self.path)
 453
 454        ctes = outputs.get("ctes")
 455        query = outputs.get("query")
 456        partial = outputs.pop("partial", None)
 457
 458        if ctes is None and query is None:
 459            _raise_error("Incomplete test, outputs must contain 'query' or 'ctes'", self.path)
 460
 461        def _normalize_rows(
 462            values: t.List[Row] | t.Dict,
 463            name: str,
 464            partial: bool = False,
 465            dialect: DialectType = None,
 466        ) -> t.Dict:
 467            import pandas as pd
 468
 469            if not isinstance(values, dict):
 470                values = {"rows": values}
 471
 472            rows = values.get("rows")
 473            query = values.get("query")
 474
 475            fmt = values.get("format")
 476            path = values.get("path")
 477            if fmt == "csv":
 478                csv_settings = values.get("csv_settings") or {}
 479                rows = pd.read_csv(path or StringIO(rows), **csv_settings).to_dict(orient="records")
 480            elif fmt in (None, "yaml"):
 481                if path:
 482                    input_rows = yaml_load(Path(path))
 483                    rows = input_rows.get("rows") if isinstance(input_rows, dict) else input_rows
 484            else:
 485                _raise_error(f"Unsupported data format '{fmt}' for '{name}'", self.path)
 486
 487            if query is not None:
 488                if rows is not None:
 489                    _raise_error(
 490                        f"Invalid test, cannot set both 'query' and 'rows' for '{name}'", self.path
 491                    )
 492
 493                # We parse the user-supplied query using the testing adapter dialect, but we
 494                # normalize its identifiers according to the model's dialect, so that, e.g.,
 495                # the projection names match those in its `columns_to_types` field
 496                values["query"] = normalize_identifiers(
 497                    exp.maybe_parse(query, dialect=self._test_adapter_dialect), dialect=dialect
 498                )
 499                return values
 500
 501            if rows is None:
 502                _raise_error(f"Incomplete test, missing row data for '{name}'", self.path)
 503
 504            assert isinstance(rows, list)
 505            values["rows"] = [
 506                {self._normalize_column_name(column): value for column, value in row.items()}
 507                for row in rows
 508            ]
 509            if partial:
 510                values["partial"] = True
 511
 512            return values
 513
 514        def _normalize_sources(
 515            sources: t.Dict, partial: bool = False, with_default_catalog: bool = True
 516        ) -> t.Dict:
 517            normalized_sources = {}
 518            for name, values in sources.items():
 519                normalized_name = self._normalize_model_name(
 520                    name, with_default_catalog=with_default_catalog
 521                )
 522                model = self.models.get(normalized_name)
 523                dialect = model.dialect if model else self.dialect
 524
 525                normalized_sources[normalized_name] = _normalize_rows(
 526                    values, name, partial=partial, dialect=dialect
 527                )
 528
 529            return normalized_sources
 530
 531        normalized_model_name = self._normalize_model_name(self.body["model"])
 532        self.body["model"] = normalized_model_name
 533
 534        if inputs:
 535            inputs = _normalize_sources(inputs)
 536            for name, values in inputs.items():
 537                columns = values.get("columns")
 538                if columns is None:
 539                    continue
 540
 541                if not isinstance(columns, dict):
 542                    _raise_error(
 543                        f"Invalid 'columns' value for model '{name}', expected a mapping name -> type",
 544                        self.path,
 545                    )
 546
 547                values["columns"] = {
 548                    self._normalize_column_name(c): exp.DataType.build(
 549                        t, dialect=self._test_adapter_dialect
 550                    )
 551                    for c, t in columns.items()
 552                }
 553
 554            for depends_on in self.model.depends_on:
 555                if depends_on not in inputs:
 556                    _raise_error(f"Incomplete test, missing input model '{depends_on}'", self.path)
 557
 558            if self.model.depends_on_self and normalized_model_name not in inputs:
 559                inputs[normalized_model_name] = {"rows": []}
 560
 561            self.body["inputs"] = inputs
 562
 563        if ctes:
 564            outputs["ctes"] = _normalize_sources(ctes, partial=partial, with_default_catalog=False)
 565
 566        if query or query == []:
 567            outputs["query"] = _normalize_rows(
 568                query, self.model.name, partial=partial, dialect=self.model.dialect
 569            )
 570
 571    def _test_fixture_table(self, name: str) -> exp.Table:
 572        table = self._fixture_table_cache.get(name)
 573        if not table:
 574            table = exp.to_table(name, dialect=self._test_adapter_dialect)
 575
 576            # We change the table path below, so this ensures there are no name clashes
 577            table.this.set("this", "__".join(part.name for part in table.parts))
 578
 579            table.set("db", self._fixture_schema.copy())
 580            if self._fixture_catalog:
 581                table.set("catalog", self._fixture_catalog.copy())
 582
 583            self._fixture_table_cache[name] = table
 584
 585        return table
 586
 587    def _normalize_model_name(self, name: str, with_default_catalog: bool = True) -> str:
 588        normalized_name = self._normalized_model_name_cache.get((name, with_default_catalog))
 589        if normalized_name is None:
 590            default_catalog = self.default_catalog if with_default_catalog else None
 591            normalized_name = normalize_model_name(
 592                name, default_catalog=default_catalog, dialect=self.dialect
 593            )
 594            self._normalized_model_name_cache[(name, with_default_catalog)] = normalized_name
 595
 596        return normalized_name
 597
 598    def _normalize_column_name(self, name: str) -> str:
 599        normalized_name = self._normalized_column_name_cache.get(name)
 600        if normalized_name is None:
 601            normalized_name = normalize_identifiers(name, dialect=self.dialect).name
 602            self._normalized_column_name_cache[name] = normalized_name
 603
 604        return normalized_name
 605
 606    @contextmanager
 607    def _concurrent_render_context(self) -> t.Iterator[None]:
 608        """
 609        Context manager that ensures that the tests are executed safely in a concurrent environment.
 610        This is needed in case `execution_time` is set, as we'd then have to:
 611        - Freeze time through `time_machine` (not thread safe)
 612        - Globally patch the SQLGlot dialect so that any date/time nodes are evaluated at the `execution_time` during generation
 613        """
 614        import time_machine
 615
 616        lock_ctx: AbstractContextManager = (
 617            self.CONCURRENT_RENDER_LOCK if self.concurrency else nullcontext()
 618        )
 619        time_ctx: AbstractContextManager = nullcontext()
 620        dialect_patch_ctx: AbstractContextManager = nullcontext()
 621
 622        if self._execution_time:
 623            time_ctx = time_machine.travel(self._execution_time, tick=False)
 624            dialect_patch_ctx = patch.dict(
 625                self._test_adapter_dialect.generator_class.TRANSFORMS, self._transforms
 626            )
 627
 628        with lock_ctx, time_ctx, dialect_patch_ctx:
 629            yield
 630
 631    def _execute(self, query: exp.Query | str) -> pd.DataFrame:
 632        """Executes the given query using the testing engine adapter and returns a DataFrame."""
 633        return self.engine_adapter.fetchdf(query)
 634
 635    def _create_df(
 636        self,
 637        values: t.Dict[str, t.Any],
 638        columns: t.Optional[t.Collection] = None,
 639        partial: t.Optional[bool] = False,
 640    ) -> pd.DataFrame:
 641        import pandas as pd
 642
 643        query = values.get("query")
 644        if query:
 645            if not partial:
 646                query = self._add_missing_columns(query, columns)
 647
 648            return self._execute(query)
 649
 650        rows = values["rows"]
 651        columns_str: t.Optional[t.List[str]] = None
 652        if columns:
 653            columns_str = [str(c) for c in columns]
 654            referenced_columns = list(dict.fromkeys(col for row in rows for col in row))
 655            _raise_if_unexpected_columns(columns, referenced_columns)
 656
 657            if partial:
 658                columns_str = [c for c in columns_str if c in referenced_columns]
 659
 660        return pd.DataFrame.from_records(rows, columns=columns_str)
 661
 662    def _add_missing_columns(
 663        self, query: exp.Query, all_columns: t.Optional[t.Collection[str]] = None
 664    ) -> exp.Query:
 665        if not all_columns or query.is_star:
 666            return query
 667
 668        query_columns = set(query.named_selects)
 669        missing_columns = [col for col in all_columns if col not in query_columns]
 670        if missing_columns:
 671            query.select(*[exp.null().as_(col) for col in missing_columns], copy=False)
 672
 673        return query
 674
 675
 676class SqlModelTest(ModelTest):
 677    def test_ctes(self, ctes: t.Dict[str, exp.Expr], recursive: bool = False) -> None:
 678        """Run CTE queries and compare output to expected output"""
 679        for cte_name, values in self.body["outputs"].get("ctes", {}).items():
 680            with self.subTest(cte=cte_name):
 681                if cte_name not in ctes:
 682                    _raise_error(
 683                        f"No CTE named {cte_name} found in model {self.model.name}", self.path
 684                    )
 685
 686                cte_query = ctes[cte_name].this
 687
 688                sort = cte_query.args.get("order") is None
 689                partial = values.get("partial")
 690
 691                cte_query = exp.select(*_projection_identifiers(cte_query)).from_(cte_name)
 692                for alias, cte in ctes.items():
 693                    cte_query = cte_query.with_(alias, cte.this, recursive=recursive)
 694
 695                with self._concurrent_render_context():
 696                    # Similar to the model's query, we render the CTE query under the locked context
 697                    # so that the execution (fetchdf) can continue concurrently between the threads
 698                    sql = cte_query.sql(
 699                        self._test_adapter_dialect, pretty=self.engine_adapter._pretty_sql
 700                    )
 701
 702                actual = self._execute(sql)
 703                expected = self._create_df(values, columns=cte_query.named_selects, partial=partial)
 704
 705                self.assert_equal(expected, actual, sort=sort, partial=partial)
 706
 707    def runTest(self) -> None:
 708        with self._concurrent_render_context():
 709            # Render the model's query and generate the SQL under the locked context so that
 710            # execution (fetchdf) can continue concurrently between the threads
 711            query = self._render_model_query()
 712            sql = query.sql(self._test_adapter_dialect, pretty=self.engine_adapter._pretty_sql)
 713
 714        with_clause = query.args.get("with_")
 715
 716        if with_clause:
 717            self.test_ctes(
 718                {
 719                    self._normalize_model_name(cte.alias, with_default_catalog=False): cte
 720                    for cte in query.ctes
 721                },
 722                recursive=with_clause.recursive,
 723            )
 724
 725        values = self.body["outputs"].get("query")
 726        if values is not None:
 727            partial = values.get("partial")
 728            sort = query.args.get("order") is None
 729
 730            actual = self._execute(sql)
 731            expected = self._create_df(values, columns=self.model.columns_to_types, partial=partial)
 732
 733            self.assert_equal(expected, actual, sort=sort, partial=partial)
 734
 735    def _render_model_query(self) -> exp.Query:
 736        variables = self.body.get("vars", {}).copy()
 737        time_kwargs = {key: variables.pop(key) for key in TIME_KWARG_KEYS if key in variables}
 738
 739        query = self.model.render_query_or_raise(
 740            **time_kwargs,
 741            variables=variables,
 742            engine_adapter=self.engine_adapter,
 743            table_mapping={
 744                name: self._test_fixture_table(name).sql() for name in self.body.get("inputs", {})
 745            },
 746            runtime_stage=RuntimeStage.TESTING,
 747        )
 748        return query
 749
 750
 751class PythonModelTest(ModelTest):
 752    def __init__(
 753        self,
 754        body: t.Dict[str, t.Any],
 755        test_name: str,
 756        model: Model,
 757        models: UniqueKeyDict[str, Model],
 758        engine_adapter: EngineAdapter,
 759        dialect: str | None = None,
 760        path: Path | None = None,
 761        preserve_fixtures: bool = False,
 762        default_catalog: str | None = None,
 763        concurrency: bool = False,
 764        verbosity: Verbosity = Verbosity.DEFAULT,
 765    ) -> None:
 766        """PythonModelTest encapsulates a unit test for a Python model.
 767
 768        Args:
 769            body: A dictionary that contains test metadata like inputs and outputs.
 770            test_name: The name of the test.
 771            model: The Python model that is being tested.
 772            models: All models to use for expansion and mapping of physical locations.
 773            engine_adapter: The engine adapter to use.
 774            dialect: The models' dialect, used for normalization purposes.
 775            path: An optional path to the test definition yaml file.
 776            preserve_fixtures: Preserve the fixture tables in the testing database, useful for debugging.
 777        """
 778        from sqlmesh.core.test.context import TestExecutionContext
 779
 780        super().__init__(
 781            body,
 782            test_name,
 783            model,
 784            models,
 785            engine_adapter,
 786            dialect,
 787            path,
 788            preserve_fixtures,
 789            default_catalog,
 790            concurrency,
 791            verbosity,
 792        )
 793
 794        self.context = TestExecutionContext(
 795            engine_adapter=engine_adapter,
 796            models=models,
 797            test=self,
 798            default_dialect=dialect,
 799            default_catalog=default_catalog,
 800        )
 801
 802    def runTest(self) -> None:
 803        values = self.body["outputs"].get("query")
 804        if values is not None:
 805            partial = values.get("partial")
 806
 807            actual_df = self._execute_model()
 808            actual_df.reset_index(drop=True, inplace=True)
 809            expected = self._create_df(values, columns=self.model.columns_to_types, partial=partial)
 810
 811            self.assert_equal(expected, actual_df, sort=True, partial=partial)
 812
 813    def _execute_model(self) -> pd.DataFrame:
 814        """Executes the python model and returns a DataFrame."""
 815        import pandas as pd
 816
 817        with self._concurrent_render_context():
 818            variables = self.body.get("vars", {}).copy()
 819            time_kwargs = {key: variables.pop(key) for key in TIME_KWARG_KEYS if key in variables}
 820            df = next(self.model.render(context=self.context, variables=variables, **time_kwargs))
 821
 822        assert not isinstance(df, exp.Expr)
 823        return df if isinstance(df, pd.DataFrame) else df.toPandas()
 824
 825
 826def generate_test(
 827    model: Model,
 828    input_queries: t.Dict[str, str],
 829    models: UniqueKeyDict[str, Model],
 830    engine_adapter: EngineAdapter,
 831    test_engine_adapter: EngineAdapter,
 832    project_path: Path,
 833    overwrite: bool = False,
 834    variables: t.Optional[t.Dict[str, str]] = None,
 835    path: t.Optional[str] = None,
 836    name: t.Optional[str] = None,
 837    include_ctes: bool = False,
 838) -> None:
 839    """Generate a unit test fixture for a given model.
 840
 841    Args:
 842        model: The model to test.
 843        input_queries: Mapping of model names to queries. Each model included in this mapping
 844            will be populated in the test based on the results of the corresponding query.
 845        models: The context's models.
 846        engine_adapter: The target engine adapter.
 847        test_engine_adapter: The test engine adapter.
 848        project_path: The path pointing to the project's root directory.
 849        overwrite: Whether to overwrite the existing test in case of a file path collision.
 850            When set to False, an error will be raised if there is such a collision.
 851        variables: Key-value pairs that will define variables needed by the model.
 852        path: The file path corresponding to the fixture, relative to the test directory.
 853            By default, the fixture will be created under the test directory and the file name
 854            will be inferred from the test's name.
 855        name: The name of the test. This is inferred from the model name by default.
 856        include_ctes: When true, CTE fixtures will also be generated.
 857    """
 858    import numpy as np
 859
 860    test_name = name or f"test_{model.view_name}"
 861    path = path or f"{test_name}.yaml"
 862
 863    extension = path.split(".")[-1].lower()
 864    if extension not in ("yaml", "yml"):
 865        path = f"{path}.yaml"
 866
 867    fixture_path = project_path / c.TESTS / path
 868    if not overwrite and fixture_path.exists():
 869        raise ConfigError(
 870            f"Fixture '{fixture_path}' already exists, make sure to set --overwrite if it can be safely overwritten."
 871        )
 872
 873    # ruamel.yaml does not support pandas Timestamps, so we must convert them to python
 874    # datetime or datetime.date objects based on column type
 875    inputs = {
 876        dep: pandas_timestamp_to_pydatetime(
 877            engine_adapter.fetchdf(query).apply(lambda col: col.map(_normalize_df_value)),
 878            models[dep].columns_to_types,
 879        )
 880        .replace({np.nan: None})
 881        .to_dict(orient="records")
 882        for dep, query in input_queries.items()
 883    }
 884    outputs: t.Dict[str, t.Any] = {"query": {}}
 885    variables = variables or {}
 886    test_body = {"model": model.fqn, "inputs": inputs, "outputs": outputs}
 887
 888    if variables:
 889        test_body["vars"] = variables
 890
 891    test = ModelTest.create_test(
 892        body=test_body.copy(),
 893        test_name=test_name,
 894        models=models,
 895        engine_adapter=test_engine_adapter,
 896        dialect=model.dialect,
 897        path=fixture_path,
 898        default_catalog=model.default_catalog,
 899    )
 900    if not test:
 901        return
 902
 903    test.setUp()
 904
 905    if isinstance(model, SqlModel):
 906        assert isinstance(test, SqlModelTest)
 907        model_query = test._render_model_query()
 908        with_clause = model_query.args.get("with_")
 909
 910        if with_clause and include_ctes:
 911            ctes = {}
 912            recursive = with_clause.recursive
 913            previous_ctes: t.List[exp.CTE] = []
 914
 915            for cte in model_query.ctes:
 916                cte_query = cte.this
 917                cte_identifier = cte.args["alias"].this
 918
 919                cte_query = exp.select(*_projection_identifiers(cte_query)).from_(cte_identifier)
 920
 921                for prev in chain(previous_ctes, [cte]):
 922                    cte_query = cte_query.with_(
 923                        prev.args["alias"].this, prev.this, recursive=recursive
 924                    )
 925
 926                cte_output = test._execute(cte_query)
 927                ctes[cte.alias] = (
 928                    pandas_timestamp_to_pydatetime(
 929                        df=cte_output.apply(lambda col: col.map(_normalize_df_value)),
 930                    )
 931                    .replace({np.nan: None})
 932                    .to_dict(orient="records")
 933                )
 934
 935                previous_ctes.append(cte)
 936
 937            if ctes:
 938                outputs["ctes"] = ctes
 939
 940        output = test._execute(model_query)
 941    else:
 942        output = t.cast(PythonModelTest, test)._execute_model()
 943
 944    outputs["query"] = (
 945        pandas_timestamp_to_pydatetime(
 946            output.apply(lambda col: col.map(_normalize_df_value)), model.columns_to_types
 947        )
 948        .replace({np.nan: None})
 949        .to_dict(orient="records")
 950    )
 951
 952    test.tearDown()
 953
 954    fixture_path.parent.mkdir(exist_ok=True, parents=True)
 955    with open(fixture_path, "w", encoding="utf-8") as file:
 956        yaml.dump({test_name: test_body}, file)
 957
 958
 959def _projection_identifiers(query: exp.Query) -> t.List[str | exp.Identifier]:
 960    identifiers: t.List[str | exp.Identifier] = []
 961    for select in query.selects:
 962        if isinstance(select, exp.Alias):
 963            identifiers.append(select.args["alias"])
 964        elif isinstance(select, exp.Column):
 965            identifiers.append(select.this)
 966        else:
 967            identifiers.append(select.output_name)
 968
 969    return identifiers
 970
 971
 972def _raise_if_unexpected_columns(
 973    expected_cols: t.Collection[str], actual_cols: t.Collection[str]
 974) -> None:
 975    unique_expected_cols = set(expected_cols)
 976    unknown_cols = [col for col in actual_cols if col not in unique_expected_cols]
 977
 978    if unknown_cols:
 979        expected = f"Expected column(s): {', '.join(list(expected_cols))}\n"
 980        unknown = f"Unknown column(s): {', '.join(unknown_cols)}"
 981        _raise_error(f"Detected unknown column(s)\n\n{expected}{unknown}")
 982
 983
 984def _row_difference(left: pd.DataFrame, right: pd.DataFrame) -> pd.DataFrame:
 985    """Returns all rows in `left` that don't appear in `right`."""
 986    import numpy as np
 987    import pandas as pd
 988
 989    rows_missing_from_right = []
 990
 991    # `None` replaces `np.nan` because `np.nan != np.nan` and this would affect the mapping lookup
 992    right_row_count: t.MutableMapping[t.Tuple, int] = Counter(
 993        right.replace({np.nan: None}).itertuples(index=False, name=None)
 994    )
 995    for left_row in left.replace({np.nan: None}).itertuples(index=False):
 996        left_row_tuple = tuple(left_row)
 997        if right_row_count[left_row_tuple] <= 0:
 998            rows_missing_from_right.append(left_row)
 999        else:
1000            right_row_count[left_row_tuple] -= 1
1001
1002    return pd.DataFrame(rows_missing_from_right)
1003
1004
1005def _raise_error(msg: str, path: Path | None = None) -> None:
1006    if path:
1007        raise TestError(f"Failed to run test at {path}:\n{msg}")
1008    raise TestError(f"Failed to run test:\n{msg}")
1009
1010
1011def _normalize_df_value(value: t.Any) -> t.Any:
1012    """Normalize data in a pandas dataframe so ruamel and sqlglot can deal with it."""
1013    import numpy as np
1014
1015    if isinstance(value, (list, np.ndarray)):
1016        return [_normalize_df_value(v) for v in value]
1017    if isinstance(value, dict):
1018        if "key" in value and "value" in value:
1019            # Maps returned by DuckDB look like: {'key': ['key1', 'key2'], 'value': [10, 20]}
1020            # so we convert to {'key1': 10, 'key2': 20} (TODO: handle more dialects here)
1021            return {k: _normalize_df_value(v) for k, v in zip(value["key"], value["value"])}
1022        return {k: _normalize_df_value(v) for k, v in value.items()}
1023    return value
1024
1025
1026def _split_df_by_column_pairs(df: pd.DataFrame, pairs_per_chunk: int = 4) -> t.List[pd.DataFrame]:
1027    """Split a dataframe into chunks of column pairs.
1028
1029    Args:
1030        df: The dataframe to split
1031        pairs_per_chunk: Number of column pairs per chunk (default: 4)
1032
1033    Returns:
1034        List of dataframes, each containing an even number of columns
1035    """
1036    total_columns = len(df.columns)
1037
1038    # If we have fewer columns than pairs_per_chunk * 2, return the original df
1039    if total_columns <= pairs_per_chunk * 2:
1040        return [df]
1041
1042    # Calculate number of chunks needed to split columns evenly
1043    num_chunks = (total_columns + (pairs_per_chunk * 2 - 1)) // (pairs_per_chunk * 2)
1044
1045    # Calculate columns per chunk to ensure equal distribution
1046    # We round down to nearest even number to ensure each chunk has even columns
1047    columns_per_chunk = (total_columns // num_chunks) & ~1  # Round down to nearest even number
1048    remainder = total_columns - (columns_per_chunk * num_chunks)
1049
1050    chunks = []
1051    start_idx = 0
1052
1053    # Distribute columns evenly across chunks
1054    for i in range(num_chunks):
1055        # Add 2 columns to early chunks if there's a remainder
1056        # This ensures we always add pairs of columns
1057        extra = 2 if i < remainder // 2 else 0
1058        end_idx = start_idx + columns_per_chunk + extra
1059        chunk = df.iloc[:, start_idx:end_idx]
1060        chunks.append(chunk)
1061        start_idx = end_idx
1062
1063    return chunks