feature: adding smarter unit test execution (#5966)

Signed-off-by: Cortland Goffena <30168413+cmgoffena13@users.noreply.github.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

Author: cmgoffena13

docs/concepts/overview.md

@@ -52,7 +52,9 @@ Bad data is worse than no data. The best way to keep bad data out of your system
 ### [Tests](./tests.md)
 SQLMesh "tests" are similar to unit tests in software development, where the unit is a single model. SQLMesh tests validate model *code* — you specify the input data and expected output, then SQLMesh runs the test and compares the expected and actual output.
 
-SQLMesh automatically runs tests when you apply a `plan`, or you can run them on demand with the [`test` command](../reference/cli.md#test).
+SQLMesh automatically runs all unit tests when a `plan` is created. Use `--test-changed-only` to run tests only for models included in the plan (added, modified, or restated), `--skip-tests` to skip, run tests for specific models with [`sqlmesh test --select-model`](../reference/cli.md#test), or run the full suite on demand with the [`test` command](../reference/cli.md#test).
+
+Learn more in the [testing guide](../guides/testing.md).
 
 ### [Audits](./audits.md)
 In contrast to tests, SQLMesh "audits" validate the results of model code applied to your actual data.

docs/concepts/tests.md

@@ -1,6 +1,8 @@
 # Testing
 
-Testing allows you to protect your project from regression by continuously verifying that the output of each model matches your expectations. Unlike [audits](audits.md), tests are executed either on demand (for example, as part of a CI/CD job) or every time a new [plan](plans.md) is created.
+Testing allows you to protect your project from regression by continuously verifying that the output of each model matches your expectations. Unlike [audits](audits.md), tests are executed either on demand (for example, as part of a CI/CD job or via [`sqlmesh test`](../reference/cli.md#test)) or when a new [plan](plans.md) is created.
+
+By default, `sqlmesh plan` runs all unit tests. Use `--test-changed-only` to run tests only for models included in the plan (added, modified, or restated), or `--skip-tests` to run none. With both `--select-model` and `--test-changed-only`, tests run only for selected models that changed.
 
 Similar to unit testing in software development, SQLMesh evaluates the model's logic against predefined inputs and then compares the output to expected outcomes provided as part of each test.
 

docs/faq/faq.md

@@ -102,7 +102,7 @@
     SQLMesh's default behavior is appropriate for most deployments, but you can override where SQLMesh creates physical tables and views with [schema configuration options](../guides/configuration.md#environment-schemas).
 
 ??? question "What's the difference between a `test` and an `audit`?"
-    A SQLMesh [`test`](../concepts/tests.md) is analogous to a "unit test" in software engineering. It tests *code* based on known inputs and outputs. In SQLMesh, the inputs and outputs are specified in a YAML file, and SQLMesh automatically runs them when `sqlmesh plan` is executed.
+    A SQLMesh [`test`](../concepts/tests.md) is analogous to a "unit test" in software engineering. It tests *code* based on known inputs and outputs. In SQLMesh, the inputs and outputs are specified in a YAML file, and SQLMesh runs all unit tests when `sqlmesh plan` is executed (use `--test-changed-only` to run only tests for models in the plan).
 
     Writing YAML is annoying and error-prone, so SQLMesh's [`create_test` command](../concepts/tests.md#automatic-test-generation) allows you to automatically generate YAML test files based on queries of existing data tables.
 

docs/guides/models.md

@@ -161,7 +161,7 @@ Reverting to a previous model version is a quick operation since no additional w
 
 SQLMesh automatically validates your models in order to ensure the quality and accuracy of your data. This is done via the following:
 
-* Running unit tests by default when you execute the `plan` command. This ensures all changes to applied to any environment are logically validated. Refer to [testing](../concepts/tests.md) for more information.
+* Running all unit tests when you execute the `plan` command (use `--test-changed-only` to run only tests for models in the plan). This ensures changes applied to any environment are logically validated. Refer to [testing](../concepts/tests.md) for more information.
 * Running audits whenever data is loaded to a table (either for backfill or loading on a cadence). This way you know all data present in any table has passed all defined audits. Refer to [auditing](../concepts/audits.md) for more information.
 
 SQLMesh also provides automatic validation via CI/CD by automatically creating a preview environment.

docs/guides/testing.md

@@ -12,11 +12,19 @@ OK
 ```
 As the unit tests run, SQLMesh will identify any that fail.
 
+By default, `sqlmesh plan` runs all unit tests. Use `--test-changed-only` to run tests only for models included in the plan (added, modified, or restated), or `--skip-tests` to run none. With both `--select-model` and `--test-changed-only`, tests run only for selected models that changed.
+
 For more information about tests, refer to [testing](../concepts/tests.md).
 
 ### Test changes to a specific model
 
-To run a specific model test, pass in the suite file name followed by `::` and the name of the test; for example: `sqlmesh test tests/test_suite.yaml::test_example_full_model`.
+To run unit tests for a specific model, use `--select-model`:
+
+```bash
+$ sqlmesh test --select-model sqlmesh_example.full_model
+```
+
+Alternatively, pass in the suite file name followed by `::` and the name of the test; for example: `sqlmesh test tests/test_suite.yaml::test_example_full_model`.
 
 ### Run a subset of tests
 

docs/reference/cli.md

@@ -357,6 +357,8 @@ Options:
                                   Default: prod.
   --skip-tests                    Skip tests prior to generating the plan if
                                   they are defined.
+  --test-changed-only             Run unit tests only for models included in
+                                  the plan instead of all tests.
   --skip-linter                   Skip linting prior to generating the plan if
                                   the linter is enabled.
   -r, --restate-model TEXT        Restate data for specified models and models
@@ -626,6 +628,8 @@ Options:
   -v, --verbose        Verbose output.
   --preserve-fixtures  Preserve the fixture tables in the testing database,
                        useful for debugging.
+  --select-model TEXT  Select specific models to run unit tests for. Can be
+                       specified multiple times.
   --help               Show this message and exit.
 ```
 

docs/reference/notebook.md

@@ -96,7 +96,7 @@ options:
 #### plan
 ```
 %plan [--start START] [--end END] [--execution-time EXECUTION_TIME]
-            [--create-from CREATE_FROM] [--skip-tests]
+            [--create-from CREATE_FROM] [--skip-tests] [--test-changed-only]
             [--restate-model [RESTATE_MODEL ...]] [--no-gaps]
             [--skip-backfill, --dry-run] [--forward-only]
             [--effective-from EFFECTIVE_FROM] [--no-prompts] [--auto-apply]
@@ -120,6 +120,8 @@ options:
                         The environment to create the target environment from
                         if it doesn't exist. Default: prod.
   --skip-tests, -t      Skip the unit tests defined for the model.
+  --test-changed-only   Run unit tests only for models included in the plan
+                        instead of all tests.
   --restate-model <[RESTATE_MODEL ...]>, -r <[RESTATE_MODEL ...]>
                         Restate data for specified models (and models
                         downstream from the one specified). For production
@@ -131,7 +133,8 @@ options:
                         comparing to existing snapshots for matching models in
                         the target environment.
   --skip-backfill, --dry-run
-                        Skip the backfill step and only create a virtual update for the plan.
+                        Skip the backfill step and only create a virtual
+                        update for the plan.
   --forward-only        Create a plan for forward-only changes.
   --effective-from EFFECTIVE_FROM
                         The effective date from which to apply forward-only
@@ -430,7 +433,8 @@ options:
 
 #### run_test
 ```
-%run_test [--pattern [PATTERN ...]] [--verbose] [--preserve-fixtures] [tests ...]
+%run_test [--pattern [PATTERN ...]] [--verbose] [--preserve-fixtures]
+          [--select-model [SELECT_MODEL ...]] [tests ...]
 
 Run unit test(s).
 
@@ -443,6 +447,8 @@ options:
   --verbose, -v         Verbose output.
   --preserve-fixtures   Preserve the fixture tables in the testing database,
                         useful for debugging.
+  --select-model <[SELECT_MODEL ...]>
+                        Select specific models to run unit tests for.
 ```
 
 #### audit

sqlmesh/cli/main.py

@@ -427,6 +427,12 @@ def diff(ctx: click.Context, environment: t.Optional[str] = None) -> None:
     help="Skip tests prior to generating the plan if they are defined.",
     default=None,
 )
+@click.option(
+    "--test-changed-only",
+    is_flag=True,
+    help="Run unit tests only for models included in the plan instead of all tests.",
+    default=None,
+)
 @click.option(
     "--skip-linter",
     is_flag=True,
@@ -799,6 +805,12 @@ def create_test(
     default=False,
     help="Preserve the fixture tables in the testing database, useful for debugging.",
 )
+@click.option(
+    "--select-model",
+    type=str,
+    multiple=True,
+    help="Select specific models to run unit tests for.",
+)
 @click.argument("tests", nargs=-1)
 @click.pass_obj
 @error_handler
@@ -808,14 +820,19 @@ def test(
     k: t.List[str],
     verbose: int,
     preserve_fixtures: bool,
+    select_model: t.List[str],
     tests: t.List[str],
 ) -> None:
     """Run model unit tests."""
+    model_names = (
+        obj._new_selector().expand_model_selections(select_model) if select_model else None
+    )
     result = obj.test(
         match_patterns=k,
         tests=tests,
         verbosity=Verbosity(verbose),
         preserve_fixtures=preserve_fixtures,
+        model_names=model_names,
     )
     if not result.wasSuccessful():
         exit(1)

sqlmesh/core/console.py

@@ -2232,6 +2232,9 @@ def log_test_results(self, result: ModelTextTestResult, target_dialect: str) ->
         message = (
             f"Ran {result.testsRun} tests against {target_dialect} in {result.duration} seconds."
         )
+        if result.tests_skipped:
+            message = f"{message}\nSkipped {result.tests_skipped} tests"
+
         if result.wasSuccessful():
             self._print("=" * divider_length)
             self._print(
@@ -3160,6 +3163,8 @@ def log_test_results(self, result: ModelTextTestResult, target_dialect: str) ->
         message = (
             f"Ran {result.testsRun} tests against {target_dialect} in {result.duration} seconds."
         )
+        if result.tests_skipped:
+            message = f"{message}\nSkipped {result.tests_skipped} tests"
 
         if result.wasSuccessful():
             success_color = {"color": "#008000"}
@@ -3600,6 +3605,8 @@ def log_test_results(self, result: ModelTextTestResult, target_dialect: str) ->
             return
 
         message = f"Ran `{result.testsRun}` Tests Against `{target_dialect}`"
+        if result.tests_skipped:
+            message = f"{message}\n**Skipped `{result.tests_skipped}` Tests**"
 
         if result.wasSuccessful():
             self._print(f"**Successfully {message}**\n\n")

sqlmesh/core/context.py

@@ -115,6 +115,7 @@
     ModelTestMetadata,
     generate_test,
     run_tests,
+    filter_tests_by_model_names,
     filter_tests_by_patterns,
 )
 from sqlmesh.core.user import User
@@ -1416,6 +1417,7 @@ def plan(
         execution_time: t.Optional[TimeLike] = None,
         create_from: t.Optional[str] = None,
         skip_tests: t.Optional[bool] = None,
+        test_changed_only: t.Optional[bool] = None,
         restate_models: t.Optional[t.Iterable[str]] = None,
         no_gaps: t.Optional[bool] = None,
         skip_backfill: t.Optional[bool] = None,
@@ -1453,6 +1455,7 @@ def plan(
             create_from: The environment to create the target environment from if it
                 doesn't exist. If not specified, the "prod" environment will be used.
             skip_tests: Unit tests are run by default so this will skip them if enabled
+            test_changed_only: Run unit tests only for models included in the plan instead of all tests
             restate_models: A list of either internal or external models, or tags, that need to be restated
                 for the given plan interval. If the target environment is a production environment,
                 ALL snapshots that depended on these upstream tables will have their intervals deleted
@@ -1499,6 +1502,7 @@ def plan(
             execution_time=execution_time,
             create_from=create_from,
             skip_tests=skip_tests,
+            test_changed_only=test_changed_only,
             restate_models=restate_models,
             no_gaps=no_gaps,
             skip_backfill=skip_backfill,
@@ -1553,6 +1557,7 @@ def plan_builder(
         execution_time: t.Optional[TimeLike] = None,
         create_from: t.Optional[str] = None,
         skip_tests: t.Optional[bool] = None,
+        test_changed_only: t.Optional[bool] = None,
         restate_models: t.Optional[t.Iterable[str]] = None,
         no_gaps: t.Optional[bool] = None,
         skip_backfill: t.Optional[bool] = None,
@@ -1587,6 +1592,7 @@ def plan_builder(
             create_from: The environment to create the target environment from if it
                 doesn't exist. If not specified, the "prod" environment will be used.
             skip_tests: Unit tests are run by default so this will skip them if enabled
+            test_changed_only: Run unit tests only for models included in the plan instead of all tests
             restate_models: A list of either internal or external models, or tags, that need to be restated
                 for the given plan interval. If the target environment is a production environment,
                 ALL snapshots that depended on these upstream tables will have their intervals deleted
@@ -1628,6 +1634,7 @@ def plan_builder(
             "execution_time": execution_time,
             "create_from": create_from,
             "skip_tests": skip_tests,
+            "test_changed_only": test_changed_only,
             "restate_models": list(restate_models) if restate_models is not None else None,
             "no_gaps": no_gaps,
             "skip_backfill": skip_backfill,
@@ -1657,6 +1664,11 @@ def plan_builder(
         }
 
         skip_tests = explain or skip_tests or False
+        test_changed_only = test_changed_only or False
+
+        if skip_tests and test_changed_only:
+            raise PlanError("Cannot combine --skip-tests with --test-changed-only.")
+
         no_gaps = no_gaps or False
         skip_backfill = skip_backfill or False
         empty_backfill = empty_backfill or False
@@ -1683,8 +1695,6 @@ def plan_builder(
         if not skip_linter:
             self.lint_models()
 
-        self._run_plan_tests(skip_tests=skip_tests)
-
         environment_ttl = (
             self.environment_ttl if environment not in self.pinned_environments else None
         )
@@ -1767,6 +1777,19 @@ def plan_builder(
             *[s.name for s in context_diff.added],
         }
 
+        plan_test_model_names: t.Set[str] = {
+            *modified_model_names,
+            *(expanded_restate_models or set()),
+        }
+        if select_models and test_changed_only:
+            plan_test_model_names &= selected_fqns
+
+        self._run_plan_tests(
+            skip_tests=skip_tests,
+            test_changed_only=test_changed_only,
+            model_names=plan_test_model_names,
+        )
+
         if (
             is_dev
             and not include_unmodified
@@ -2383,14 +2406,23 @@ def test(
         verbosity: Verbosity = Verbosity.DEFAULT,
         preserve_fixtures: bool = False,
         stream: t.Optional[t.TextIO] = None,
+        model_names: t.Optional[t.Collection[str]] = None,
     ) -> ModelTextTestResult:
         """Discover and run model tests"""
         if verbosity >= Verbosity.VERBOSE:
             import pandas as pd
 
             pd.set_option("display.max_columns", None)
 
-        test_meta = self.select_tests(tests=tests, patterns=match_patterns)
+        baseline_meta = self.select_tests(tests=tests, patterns=match_patterns, model_names=None)
+        if model_names is not None:
+            test_meta = self.select_tests(
+                tests=tests, patterns=match_patterns, model_names=model_names
+            )
+            tests_skipped = len(baseline_meta) - len(test_meta)
+        else:
+            test_meta = baseline_meta
+            tests_skipped = 0
 
         result = run_tests(
             model_test_metadata=test_meta,
@@ -2404,6 +2436,7 @@ def test(
             default_catalog=self.default_catalog,
             default_catalog_dialect=self.config.dialect or "",
         )
+        result.tests_skipped = tests_skipped
 
         self.console.log_test_results(
             result,
@@ -2850,15 +2883,28 @@ def _run_tests(
         result = self.test(stream=test_output_io, verbosity=verbosity)
         return result, test_output_io.getvalue()
 
-    def _run_plan_tests(self, skip_tests: bool = False) -> t.Optional[ModelTextTestResult]:
-        if not skip_tests:
-            result = self.test()
-            if not result.wasSuccessful():
-                raise PlanError(
-                    "Cannot generate plan due to failing test(s). Fix test(s) and run again."
-                )
-            return result
-        return None
+    def _run_plan_tests(
+        self,
+        skip_tests: bool = False,
+        test_changed_only: bool = False,
+        model_names: t.Optional[t.Collection[str]] = None,
+    ) -> t.Optional[ModelTextTestResult]:
+        if skip_tests:
+            return None
+
+        effective_names: t.Optional[t.Set[str]] = None
+
+        if test_changed_only:
+            effective_names = set(model_names or [])
+            if not effective_names:
+                return None
+
+        result = self.test(model_names=effective_names)
+        if not result.wasSuccessful():
+            raise PlanError(
+                "Cannot generate plan due to failing test(s). Fix test(s) and run again."
+            )
+        return result
 
     def _warn_if_virtual_catalog_rematerialization(self, plan: "Plan") -> None:
         """Warn when ClickHouse models appear as new snapshots solely because a virtual catalog
@@ -3569,6 +3615,7 @@ def select_tests(
         self,
         tests: t.Optional[t.List[str]] = None,
         patterns: t.Optional[t.List[str]] = None,
+        model_names: t.Optional[t.Collection[str]] = None,
     ) -> t.List[ModelTestMetadata]:
         """Filter pre-loaded test metadata based on tests and patterns."""
 
@@ -3592,6 +3639,14 @@ def select_tests(
         if patterns:
             test_meta = filter_tests_by_patterns(test_meta, patterns)
 
+        if model_names is not None:
+            test_meta = filter_tests_by_model_names(
+                test_meta,
+                set(model_names),
+                default_catalog=self.default_catalog,
+                dialect=self.default_dialect,
+            )
+
         return test_meta