Fix: Propagate gateway variables to model tests

Author: VaggelisD

sqlmesh/core/context.py

@@ -109,8 +109,7 @@
 from sqlmesh.core.test import (
     ModelTextTestResult,
     generate_test,
-    get_all_model_tests,
-    run_model_tests,
+    load_model_tests,
     run_tests,
 )
 from sqlmesh.core.user import User
@@ -1781,47 +1780,30 @@ def test(
         if verbosity >= Verbosity.VERBOSE:
             pd.set_option("display.max_columns", None)
 
-        if tests:
-            result = run_model_tests(
-                tests=tests,
-                models=self._models,
-                config=self.config,
-                gateway=self.gateway,
-                dialect=self.default_dialect,
-                verbosity=verbosity,
-                patterns=match_patterns,
-                preserve_fixtures=preserve_fixtures,
-                stream=stream,
-                default_catalog=self.default_catalog,
-                default_catalog_dialect=self.engine_adapter.DIALECT,
-            )
-        else:
-            test_meta = []
-
-            for path, config in self.configs.items():
-                test_meta.extend(
-                    get_all_model_tests(
-                        path / c.TESTS,
-                        patterns=match_patterns,
-                        ignore_patterns=config.ignore_patterns,
-                        variables=config.variables,
-                    )
-                )
+        default_gateway = self.gateway or self.config.default_gateway_name
 
-            result = run_tests(
-                model_test_metadata=test_meta,
-                models=self._models,
-                config=self.config,
-                gateway=self.gateway,
-                dialect=self.default_dialect,
-                verbosity=verbosity,
-                preserve_fixtures=preserve_fixtures,
-                stream=stream,
-                default_catalog=self.default_catalog,
-                default_catalog_dialect=self.engine_adapter.DIALECT,
-            )
+        # Merge the root variables with the gateway's variables
+        variables = {**self.config.variables, **self.config.get_gateway(default_gateway).variables}
 
-        return result
+        test_meta = load_model_tests(
+            configs=self.configs,
+            tests=tests,
+            patterns=match_patterns,
+            variables=variables,
+        )
+
+        return run_tests(
+            model_test_metadata=test_meta,
+            models=self._models,
+            config=self.config,
+            default_gateway=default_gateway,
+            dialect=self.default_dialect,
+            verbosity=verbosity,
+            preserve_fixtures=preserve_fixtures,
+            stream=stream,
+            default_catalog=self.default_catalog,
+            default_catalog_dialect=self.engine_adapter.DIALECT,
+        )
 
     @python_api_analytics
     def audit(

sqlmesh/core/test/__init__.py

@@ -6,9 +6,7 @@
     filter_tests_by_patterns as filter_tests_by_patterns,
     get_all_model_tests as get_all_model_tests,
     load_model_test_file as load_model_test_file,
+    load_model_tests as load_model_tests,
 )
 from sqlmesh.core.test.result import ModelTextTestResult as ModelTextTestResult
-from sqlmesh.core.test.runner import (
-    run_model_tests as run_model_tests,
-    run_tests as run_tests,
-)
+from sqlmesh.core.test.runner import run_tests as run_tests

sqlmesh/core/test/discovery.py

@@ -8,10 +8,14 @@
 
 import ruamel
 
+from sqlmesh.core import constants as c
 from sqlmesh.utils import unique
 from sqlmesh.utils.pydantic import PydanticModel
 from sqlmesh.utils.yaml import load as yaml_load
 
+if t.TYPE_CHECKING:
+    from sqlmesh.core.config.loader import C
+
 
 class ModelTestMetadata(PydanticModel):
     path: pathlib.Path
@@ -113,3 +117,46 @@ def get_all_model_tests(
     if patterns:
         model_test_metadatas = filter_tests_by_patterns(model_test_metadatas, patterns)
     return model_test_metadatas
+
+
+def load_model_tests(
+    configs: dict[pathlib.Path, C],
+    tests: t.Optional[t.List[str]] = None,
+    patterns: list[str] | None = None,
+    variables: dict[str, t.Any] | None = None,
+) -> list[ModelTestMetadata]:
+    """Load model tests into a list of ModelTestMetadata which will be propagated to the test runner.
+
+    Args:
+        tests: A list of tests to load; If not specified, all tests are loaded
+        patterns: A list of patterns to match against.
+        variables: A dictionary of variables to use when loading the tests.
+        configs: A dictionary of configs to use when loading all the tests.
+    """
+    test_meta = []
+
+    if tests:
+        for test in tests:
+            filename, test_name = test.split("::", maxsplit=1) if "::" in test else (test, "")
+
+            test_file = load_model_test_file(pathlib.Path(filename), variables=variables)
+            if test_name:
+                test_meta.append(test_file[test_name])
+            else:
+                test_meta.extend(test_file.values())
+
+        if patterns:
+            test_meta = filter_tests_by_patterns(test_meta, patterns)
+
+    else:
+        for path, config in configs.items():
+            test_meta.extend(
+                get_all_model_tests(
+                    path / c.TESTS,
+                    patterns=patterns,
+                    ignore_patterns=config.ignore_patterns,
+                    variables=variables,
+                )
+            )
+
+    return test_meta

sqlmesh/core/test/runner.py

@@ -2,7 +2,6 @@
 
 import sys
 import time
-import pathlib
 import threading
 import typing as t
 import unittest
@@ -86,7 +85,7 @@ def run_tests(
     model_test_metadata: list[ModelTestMetadata],
     models: UniqueKeyDict[str, Model],
     config: C,
-    gateway: t.Optional[str] = None,
+    default_gateway: str,
     dialect: str | None = None,
     verbosity: Verbosity = Verbosity.DEFAULT,
     preserve_fixtures: bool = False,
@@ -102,8 +101,6 @@ def run_tests(
         verbosity: The verbosity level.
         preserve_fixtures: Preserve the fixture tables in the testing database, useful for debugging.
     """
-    default_gateway = gateway or config.default_gateway_name
-
     default_test_connection = config.get_test_connection(
         gateway_name=default_gateway,
         default_catalog=default_catalog,
@@ -128,34 +125,40 @@ def run_tests(
 
     def _run_single_test(
         metadata: ModelTestMetadata, engine_adapter: EngineAdapter
-    ) -> ModelTextTestResult:
-        test = ModelTest.create_test(
-            body=metadata.body,
-            test_name=metadata.test_name,
-            models=models,
-            engine_adapter=engine_adapter,
-            dialect=dialect,
-            path=metadata.path,
-            default_catalog=default_catalog,
-            preserve_fixtures=preserve_fixtures,
-        )
+    ) -> t.Optional[ModelTextTestResult]:
+        result: t.Optional[ModelTextTestResult] = None
+        try:
+            test = ModelTest.create_test(
+                body=metadata.body,
+                test_name=metadata.test_name,
+                models=models,
+                engine_adapter=engine_adapter,
+                dialect=dialect,
+                path=metadata.path,
+                default_catalog=default_catalog,
+                preserve_fixtures=preserve_fixtures,
+            )
 
-        result = t.cast(
-            ModelTextTestResult,
-            ModelTextTestRunner().run(t.cast(unittest.TestCase, test)),
-        )
+            result = t.cast(
+                ModelTextTestResult,
+                ModelTextTestRunner().run(t.cast(unittest.TestCase, test)),
+            )
+
+            with lock:
+                if result.successes:
+                    combined_results.addSuccess(result.successes[0])
+                elif result.errors:
+                    combined_results.addError(result.original_err[0], result.original_err[1])
+                elif result.failures:
+                    combined_results.addFailure(result.original_err[0], result.original_err[1])
+                elif result.skipped:
+                    skipped_args = result.skipped[0]
+                    combined_results.addSkip(skipped_args[0], skipped_args[1])
 
-        with lock:
-            if result.successes:
-                combined_results.addSuccess(result.successes[0])
-            elif result.errors:
-                combined_results.addError(result.original_err[0], result.original_err[1])
-            elif result.failures:
-                combined_results.addFailure(result.original_err[0], result.original_err[1])
-            elif result.skipped:
-                skipped_args = result.skipped[0]
-                combined_results.addSkip(skipped_args[0], skipped_args[1])
-        return result
+                combined_results.testsRun += 1
+
+        finally:
+            return result
 
     test_results = []
 
@@ -180,57 +183,6 @@ def _run_single_test(
 
     end_time = time.perf_counter()
 
-    combined_results.testsRun = len(test_results)
-
     combined_results.log_test_report(test_duration=end_time - start_time)
 
     return combined_results
-
-
-def run_model_tests(
-    tests: list[str],
-    models: UniqueKeyDict[str, Model],
-    config: C,
-    gateway: t.Optional[str] = None,
-    dialect: str | None = None,
-    verbosity: Verbosity = Verbosity.DEFAULT,
-    patterns: list[str] | None = None,
-    preserve_fixtures: bool = False,
-    stream: t.TextIO | None = None,
-    default_catalog: t.Optional[str] = None,
-    default_catalog_dialect: str = "",
-) -> ModelTextTestResult:
-    """Load and run tests.
-
-    Args:
-        tests: A list of tests to run, e.g. [tests/test_orders.yaml::test_single_order]
-        models: All models to use for expansion and mapping of physical locations.
-        verbosity: The verbosity level.
-        patterns: A list of patterns to match against.
-        preserve_fixtures: Preserve the fixture tables in the testing database, useful for debugging.
-    """
-    loaded_tests = []
-    for test in tests:
-        filename, test_name = test.split("::", maxsplit=1) if "::" in test else (test, "")
-        path = pathlib.Path(filename)
-
-        if test_name:
-            loaded_tests.append(load_model_test_file(path, variables=config.variables)[test_name])
-        else:
-            loaded_tests.extend(load_model_test_file(path, variables=config.variables).values())
-
-    if patterns:
-        loaded_tests = filter_tests_by_patterns(loaded_tests, patterns)
-
-    return run_tests(
-        loaded_tests,
-        models,
-        config,
-        gateway=gateway,
-        dialect=dialect,
-        verbosity=verbosity,
-        preserve_fixtures=preserve_fixtures,
-        stream=stream,
-        default_catalog=default_catalog,
-        default_catalog_dialect=default_catalog_dialect,
-    )