Move test loading logic to Loader

VaggelisD · VaggelisD · commit 99e65917eb43 · 2025-04-11T14:59:02.000+03:00
diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
@@ -108,8 +108,8 @@
 from sqlmesh.core.table_diff import TableDiff
 from sqlmesh.core.test import (
     ModelTextTestResult,
+    ModelTestMetadata,
     generate_test,
-    load_model_tests,
     run_tests,
 )
 from sqlmesh.core.user import User
@@ -363,7 +363,6 @@ def __init__(
         self._excluded_requirements: t.Set[str] = set()
         self._default_catalog: t.Optional[str] = None
         self._linters: t.Dict[str, Linter] = {}
-        self._variables_by_project_gateway: t.Dict[t.Tuple[str, str], t.Dict[str, t.Any]] = {}
         self._loaded: bool = False
 
         self.path, self.config = t.cast(t.Tuple[Path, C], next(iter(self.configs.items())))
@@ -1781,12 +1780,7 @@ def test(
         if verbosity >= Verbosity.VERBOSE:
             pd.set_option("display.max_columns", None)
 
-        test_meta = load_model_tests(
-            configs=self.configs,
-            tests=tests,
-            patterns=match_patterns,
-            get_variables=self._get_variables,
-        )
+        test_meta = self._load_model_tests(tests=tests, patterns=match_patterns)
 
         return run_tests(
             model_test_metadata=test_meta,
@@ -2463,31 +2457,15 @@ def lint_models(
                 "Linter detected errors in the code. Please fix them before proceeding."
             )
 
-    def _get_variables(
-        self, config: t.Optional[C] = None, gateway_name: t.Optional[str] = None
-    ) -> t.Dict[str, t.Any]:
-        config = config or self.config
-        gateway_name = gateway_name or self.selected_gateway
-
-        key = (config.project, gateway_name)
-        if key not in self._variables_by_project_gateway:
-            try:
-                gateway = config.get_gateway(gateway_name)
-            except ConfigError:
-                from sqlmesh.core.console import get_console
+    def _load_model_tests(
+        self, tests: t.Optional[t.List[str]] = None, patterns: list[str] | None = None
+    ) -> t.List[ModelTestMetadata]:
+        model_tests = []
 
-                get_console().log_warning(
-                    f"Gateway '{gateway_name}' not found in project '{config.project}'."
-                )
-                gateway = None
-
-            self._variables_by_project_gateway[key] = {
-                **config.variables,
-                **(gateway.variables if gateway else {}),
-                c.GATEWAY: gateway_name,
-            }
+        for loader in self._loaders:
+            model_tests.extend(loader._load_model_tests(tests=tests, patterns=patterns))
 
-        return self._variables_by_project_gateway[key]
+        return model_tests
 
 
 class Context(GenericContext[Config]):
diff --git a/sqlmesh/core/loader.py b/sqlmesh/core/loader.py
@@ -2,6 +2,7 @@
 
 import abc
 import glob
+import itertools
 import linecache
 import logging
 import os
@@ -31,11 +32,13 @@
 from sqlmesh.core.model import model as model_registry
 from sqlmesh.core.model.common import make_python_env
 from sqlmesh.core.signal import signal
+from sqlmesh.core.test import ModelTestMetadata, filter_tests_by_patterns
 from sqlmesh.utils import UniqueKeyDict, sys_path
 from sqlmesh.utils.errors import ConfigError
 from sqlmesh.utils.jinja import JinjaMacroRegistry, MacroExtractor
 from sqlmesh.utils.metaprogramming import import_python_file
-from sqlmesh.utils.yaml import YAML
+from sqlmesh.utils.yaml import YAML, load as yaml_load
+
 
 if t.TYPE_CHECKING:
     from sqlmesh.core.context import GenericContext
@@ -74,6 +77,7 @@ def __init__(self, context: GenericContext, path: Path) -> None:
         self.context = context
         self.config_path = path
         self.config = self.context.configs[self.config_path]
+        self._variables_by_gateway: t.Dict[str, t.Dict[str, t.Any]] = {}
 
     def load(self) -> LoadedProject:
         """
@@ -289,6 +293,12 @@ def _load_linting_rules(self) -> RuleSet:
         """Loads user linting rules"""
         return RuleSet()
 
+    def _load_model_tests(
+        self, tests: t.Optional[t.List[str]] = None, patterns: list[str] | None = None
+    ) -> t.List[ModelTestMetadata]:
+        """Loads YAML-based model tests"""
+        return []
+
     def _glob_paths(
         self,
         path: Path,
@@ -328,7 +338,26 @@ def _track_file(self, path: Path) -> None:
         self._path_mtimes[path] = path.stat().st_mtime
 
     def _get_variables(self, gateway_name: t.Optional[str] = None) -> t.Dict[str, t.Any]:
-        return self.context._get_variables(config=self.config, gateway_name=gateway_name)
+        gateway_name = gateway_name or self.context.selected_gateway
+
+        if gateway_name not in self._variables_by_gateway:
+            try:
+                gateway = self.config.get_gateway(gateway_name)
+            except ConfigError:
+                from sqlmesh.core.console import get_console
+
+                get_console().log_warning(
+                    f"Gateway '{gateway_name}' not found in project '{self.config.project}'."
+                )
+                gateway = None
+
+            self._variables_by_gateway[gateway_name] = {
+                **self.config.variables,
+                **(gateway.variables if gateway else {}),
+                c.GATEWAY: gateway_name,
+            }
+
+        return self._variables_by_gateway[gateway_name]
 
 
 class SqlMeshLoader(Loader):
@@ -658,6 +687,53 @@ def _load_linting_rules(self) -> RuleSet:
 
         return RuleSet(user_rules.values())
 
+    def _load_model_test_file(self, path: Path) -> dict[str, ModelTestMetadata]:
+        """Load a single model test file."""
+        model_test_metadata = {}
+        contents = yaml_load(path, get_variables=self._get_variables)
+
+        for test_name, value in contents.items():
+            model_test_metadata[test_name] = ModelTestMetadata(
+                path=path, test_name=test_name, body=value
+            )
+
+        return model_test_metadata
+
+    def _load_model_tests(
+        self, tests: t.Optional[t.List[str]] = None, patterns: list[str] | None = None
+    ) -> t.List[ModelTestMetadata]:
+        """Loads YAML-based model tests"""
+        test_meta: t.List[ModelTestMetadata] = []
+
+        if tests:
+            for test in tests:
+                filename, test_name = test.split("::", maxsplit=1) if "::" in test else (test, "")
+
+                test_file = self._load_model_test_file(Path(filename))
+                if test_name:
+                    test_meta.append(test_file[test_name])
+                else:
+                    test_meta.extend(test_file.values())
+        else:
+            search_path = Path(self.config_path) / c.TESTS
+
+            for yaml_file in itertools.chain(
+                search_path.glob("**/test*.yaml"),
+                search_path.glob("**/test*.yml"),
+            ):
+                if any(
+                    yaml_file.match(ignore_pattern)
+                    for ignore_pattern in self.config.ignore_patterns or []
+                ):
+                    continue
+
+                test_meta.extend(self._load_model_test_file(yaml_file).values())
+
+        if patterns:
+            test_meta = filter_tests_by_patterns(test_meta, patterns)
+
+        return test_meta
+
     class _Cache(CacheBase):
         def __init__(self, loader: SqlMeshLoader, config_path: Path):
             self._loader = loader
diff --git a/sqlmesh/core/test/__init__.py b/sqlmesh/core/test/__init__.py
@@ -3,7 +3,7 @@
 from sqlmesh.core.test.definition import ModelTest as ModelTest, generate_test as generate_test
 from sqlmesh.core.test.discovery import (
     ModelTestMetadata as ModelTestMetadata,
-    load_model_tests as load_model_tests,
+    filter_tests_by_patterns as filter_tests_by_patterns,
 )
 from sqlmesh.core.test.result import ModelTextTestResult as ModelTextTestResult
 from sqlmesh.core.test.runner import run_tests as run_tests
diff --git a/sqlmesh/core/test/discovery.py b/sqlmesh/core/test/discovery.py
@@ -4,17 +4,11 @@
 import itertools
 import pathlib
 import typing as t
-from collections.abc import Iterator
 
 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):
@@ -30,61 +24,6 @@ def __hash__(self) -> int:
         return self.fully_qualified_test_name.__hash__()
 
 
-def load_model_test_file(
-    path: pathlib.Path,
-    config: C,
-    get_variables: t.Callable[[t.Optional[C], t.Optional[str]], t.Dict[str, str]],
-) -> dict[str, ModelTestMetadata]:
-    """Load a single model test file.
-
-    Args:
-        path: The path to the test file
-
-    returns:
-        A list of ModelTestMetadata named tuples.
-    """
-    model_test_metadata = {}
-    contents = yaml_load(path, config=config, get_variables=get_variables)
-
-    for test_name, value in contents.items():
-        model_test_metadata[test_name] = ModelTestMetadata(
-            path=path, test_name=test_name, body=value
-        )
-    return model_test_metadata
-
-
-def discover_model_tests(
-    path: pathlib.Path,
-    config: C,
-    get_variables: t.Callable[[t.Optional[C], t.Optional[str]], t.Dict[str, str]],
-) -> Iterator[ModelTestMetadata]:
-    """Discover model tests.
-
-    Model tests are defined in YAML files and contain the inputs and outputs used to test model queries.
-
-    Args:
-        path: A path to search for tests.
-        ignore_patterns: An optional list of patterns to ignore.
-
-    Returns:
-        A list of ModelTestMetadata named tuples.
-    """
-    search_path = pathlib.Path(path)
-
-    for yaml_file in itertools.chain(
-        search_path.glob("**/test*.yaml"),
-        search_path.glob("**/test*.yml"),
-    ):
-        for ignore_pattern in config.ignore_patterns or []:
-            if yaml_file.match(ignore_pattern):
-                break
-        else:
-            for model_test_metadata in load_model_test_file(
-                yaml_file, config=config, get_variables=get_variables
-            ).values():
-                yield model_test_metadata
-
-
 def filter_tests_by_patterns(
     tests: list[ModelTestMetadata], patterns: list[str]
 ) -> list[ModelTestMetadata]:
@@ -103,47 +42,3 @@ def filter_tests_by_patterns(
         if ("*" in pattern and fnmatch.fnmatchcase(test.fully_qualified_test_name, pattern))
         or pattern in test.fully_qualified_test_name
     )
-
-
-def load_model_tests(
-    configs: t.Dict[pathlib.Path, C],
-    get_variables: t.Callable[[t.Optional[C], t.Optional[str]], t.Dict[str, str]],
-    tests: t.Optional[t.List[str]] = None,
-    patterns: list[str] | 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 that'll be used to filter tests by file name.
-        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),
-                config=next(iter(configs.values())),
-                get_variables=get_variables,
-            )
-            if test_name:
-                test_meta.append(test_file[test_name])
-            else:
-                test_meta.extend(test_file.values())
-    else:
-        for path, config in configs.items():
-            test_meta.extend(
-                discover_model_tests(
-                    pathlib.Path(path / c.TESTS),
-                    config=config,
-                    get_variables=get_variables,
-                )
-            )
-
-    if patterns:
-        test_meta = filter_tests_by_patterns(test_meta, patterns)
-
-    return test_meta
diff --git a/sqlmesh/core/test/result.py b/sqlmesh/core/test/result.py
@@ -107,7 +107,6 @@ def log_test_report(self, test_duration: float) -> None:
         for _, error in errors:
             stream.writeln(unittest.TextTestResult.separator1)
             stream.writeln(f"ERROR: {error}")
-            stream.writeln(unittest.TextTestResult.separator2)
 
         # Output final report
         stream.writeln(unittest.TextTestResult.separator2)
diff --git a/sqlmesh/magics.py b/sqlmesh/magics.py
@@ -31,7 +31,7 @@
 from sqlmesh.core.context import Context
 from sqlmesh.core.dialect import format_model_expressions, parse
 from sqlmesh.core.model import load_sql_based_model
-from sqlmesh.core.test import ModelTestMetadata, load_model_tests
+from sqlmesh.core.test import ModelTestMetadata
 from sqlmesh.utils import sqlglot_dialects, yaml, Verbosity
 from sqlmesh.utils.errors import MagicError, MissingContextException, SQLMeshError
 
@@ -272,7 +272,7 @@ def test(self, context: Context, line: str, test_def_raw: t.Optional[str] = None
         if not args.test_name and not args.ls:
             raise MagicError("Must provide either test name or `--ls` to list tests")
 
-        test_meta = load_model_tests(configs=context.configs, get_variables=context._get_variables)
+        test_meta = context._load_model_tests()
 
         tests: t.Dict[str, t.Dict[str, ModelTestMetadata]] = defaultdict(dict)
         for model_test_metadata in test_meta:
diff --git a/sqlmesh/utils/yaml.py b/sqlmesh/utils/yaml.py
@@ -13,10 +13,6 @@
 from sqlmesh.utils.errors import SQLMeshError
 from sqlmesh.utils.jinja import ENVIRONMENT, create_var
 
-if t.TYPE_CHECKING:
-    from sqlmesh.core.config.loader import C
-
-
 JINJA_METHODS = {
     "env_var": lambda key, default=None: getenv(key, default),
 }
@@ -44,8 +40,7 @@ def load(
     render_jinja: bool = True,
     allow_duplicate_keys: bool = False,
     variables: t.Optional[t.Dict[str, t.Any]] = None,
-    config: t.Optional[C] = None,
-    get_variables: t.Callable[[t.Optional[C], t.Optional[str]], t.Dict[str, str]] | None = None,
+    get_variables: t.Callable[[t.Optional[str]], t.Dict[str, str]] | None = None,
 ) -> t.Dict:
     """Loads a YAML object from either a raw string or a file."""
     path: t.Optional[Path] = None
@@ -62,7 +57,7 @@ def load(
         gateway_line = GATEWAY_PATTERN.search(source)
         gateway = yaml.load(gateway_line.group(0))["gateway"] if gateway_line else None
 
-        variables = get_variables(config, gateway)
+        variables = get_variables(gateway)
 
     if render_jinja:
         source = ENVIRONMENT.from_string(source).render(
diff --git a/tests/core/test_test.py b/tests/core/test_test.py

Original file line number	Diff line number	Diff line change
`@@ -3,7 +3,7 @@`
`3`	`3`	`from sqlmesh.core.test.definition import ModelTest as ModelTest, generate_test as generate_test`
`4`	`4`	`from sqlmesh.core.test.discovery import (`
`5`	`5`	`ModelTestMetadata as ModelTestMetadata,`
`6`		`- load_model_tests as load_model_tests,`
	`6`	`+ filter_tests_by_patterns as filter_tests_by_patterns,`
`7`	`7`	`)`
`8`	`8`	`from sqlmesh.core.test.result import ModelTextTestResult as ModelTextTestResult`
`9`	`9`	`from sqlmesh.core.test.runner import run_tests as run_tests`