Repository: TobikoData/sqlmesh
Pull request #2879 (2024-07-08 20:26:31)

Feat: add support for parameterization of model names in unit tests
cc @notgerry


Patch:
diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
index 5c0fd8c6..10de8180 100644
--- a/sqlmesh/core/context.py
+++ b/sqlmesh/core/context.py
@@ -1509,6 +1509,7 @@ class GenericContext(BaseContext, t.Generic[C]):
                         path / c.TESTS,
                         patterns=match_patterns,
                         ignore_patterns=config.ignore_patterns,
+                        variables=config.variables,
                     )
                 )
 
diff --git a/sqlmesh/utils/yaml.py b/sqlmesh/utils/yaml.py
index 10cd4d97..fcf2c191 100644
--- a/sqlmesh/utils/yaml.py
+++ b/sqlmesh/utils/yaml.py
@@ -7,8 +7,9 @@ from pathlib import Path
 
 from ruamel import yaml
 
+from sqlmesh.core.constants import VAR
 from sqlmesh.utils.errors import SQLMeshError
-from sqlmesh.utils.jinja import ENVIRONMENT
+from sqlmesh.utils.jinja import ENVIRONMENT, create_var
 
 JINJA_METHODS = {
     "env_var": lambda key, default=None: getenv(key, default),
@@ -22,6 +23,7 @@ def load(
     raise_if_empty: bool = True,
     render_jinja: bool = True,
     allow_duplicate_keys: bool = False,
+    variables: t.Optional[t.Dict[str, t.Any]] = None,
 ) -> t.Dict:
     """Loads a YAML object from either a raw string or a file."""
     path: t.Optional[Path] = None
@@ -32,7 +34,12 @@ def load(
             source = file.read()
 
     if render_jinja:
-        source = ENVIRONMENT.from_string(source).render(JINJA_METHODS)
+        source = ENVIRONMENT.from_string(source).render(
+            {
+                **JINJA_METHODS,
+                VAR: create_var(variables or {}),
+            }
+        )
 
     yaml = YAML()
     yaml.allow_duplicate_keys = allow_duplicate_keys


Test patch:
diff --git a/docs/concepts/tests.md b/docs/concepts/tests.md
index d0250780..c4b2adb4 100644
--- a/docs/concepts/tests.md
+++ b/docs/concepts/tests.md
@@ -255,6 +255,24 @@ test_colors:
     execution_time: "2023-01-01 12:05:03+02:00"
 ```
 
+## Parameterized model names
+
+Testing models with parameterized names, such as `@{gold}.some_schema.some_table`, is possible using Jinja:
+
+```yaml linenums="1"
+test_parameterized_model:
+  model: {{ var('gold') }}.some_schema.some_table
+  ...
+```
+
+For example, assuming `gold` is a [config variable](../reference/configuration/#variables) with value `gold_db`, the above test would be rendered as:
+
+```yaml linenums="1"
+test_parameterized_model:
+  model: gold_db.some_schema.some_table
+  ...
+```
+
 ## Automatic test generation
 
 Creating tests manually can be repetitive and error-prone, which is why SQLMesh also provides a way to automate this process using the [`create_test` command](../reference/cli.md#create_test).
diff --git a/sqlmesh/core/test/__init__.py b/sqlmesh/core/test/__init__.py
index a8de3fd8..4b1d00de 100644
--- a/sqlmesh/core/test/__init__.py
+++ b/sqlmesh/core/test/__init__.py
@@ -111,9 +111,9 @@ def run_model_tests(
         path = pathlib.Path(filename)
 
         if test_name:
-            loaded_tests.append(load_model_test_file(path)[test_name])
+            loaded_tests.append(load_model_test_file(path, variables=config.variables)[test_name])
         else:
-            loaded_tests.extend(load_model_test_file(path).values())
+            loaded_tests.extend(load_model_test_file(path, variables=config.variables).values())
 
     if patterns:
         loaded_tests = filter_tests_by_patterns(loaded_tests, patterns)
diff --git a/sqlmesh/core/test/discovery.py b/sqlmesh/core/test/discovery.py
index a5f2f23a..a7977d20 100644
--- a/sqlmesh/core/test/discovery.py
+++ b/sqlmesh/core/test/discovery.py
@@ -26,7 +26,9 @@ class ModelTestMetadata(PydanticModel):
         return self.fully_qualified_test_name.__hash__()
 
 
-def load_model_test_file(path: pathlib.Path) -> dict[str, ModelTestMetadata]:
+def load_model_test_file(
+    path: pathlib.Path, variables: dict[str, t.Any] | None = None
+) -> dict[str, ModelTestMetadata]:
     """Load a single model test file.
 
     Args:
@@ -36,7 +38,7 @@ def load_model_test_file(path: pathlib.Path) -> dict[str, ModelTestMetadata]:
         A list of ModelTestMetadata named tuples.
     """
     model_test_metadata = {}
-    contents = yaml_load(path)
+    contents = yaml_load(path, variables=variables)
 
     for test_name, value in contents.items():
         model_test_metadata[test_name] = ModelTestMetadata(
@@ -46,7 +48,9 @@ def load_model_test_file(path: pathlib.Path) -> dict[str, ModelTestMetadata]:
 
 
 def discover_model_tests(
-    path: pathlib.Path, ignore_patterns: list[str] | None = None
+    path: pathlib.Path,
+    ignore_patterns: list[str] | None = None,
+    variables: dict[str, t.Any] | None = None,
 ) -> Iterator[ModelTestMetadata]:
     """Discover model tests.
 
@@ -69,7 +73,9 @@ def discover_model_tests(
             if yaml_file.match(ignore_pattern):
                 break
         else:
-            for model_test_metadata in load_model_test_file(yaml_file).values():
+            for model_test_metadata in load_model_test_file(
+                yaml_file, variables=variables
+            ).values():
                 yield model_test_metadata
 
 
@@ -97,9 +103,12 @@ def get_all_model_tests(
     *paths: pathlib.Path,
     patterns: list[str] | None = None,
     ignore_patterns: list[str] | None = None,
+    variables: dict[str, t.Any] | None = None,
 ) -> list[ModelTestMetadata]:
     model_test_metadatas = [
-        meta for path in paths for meta in discover_model_tests(pathlib.Path(path), ignore_patterns)
+        meta
+        for path in paths
+        for meta in discover_model_tests(pathlib.Path(path), ignore_patterns, variables=variables)
     ]
     if patterns:
         model_test_metadatas = filter_tests_by_patterns(model_test_metadatas, patterns)
diff --git a/tests/core/test_test.py b/tests/core/test_test.py
index 127bf82d..184b2b13 100644
--- a/tests/core/test_test.py
+++ b/tests/core/test_test.py
@@ -1449,6 +1449,45 @@ test_pyspark_model:
     )
 
 
+def test_variable_usage(tmp_path: Path) -> None:
+    init_example_project(tmp_path, dialect="duckdb")
+
+    config = Config(
+        default_connection=DuckDBConnectionConfig(),
+        model_defaults=ModelDefaultsConfig(dialect="duckdb"),
+        variables={"gold": "gold_db", "silver": "silver_db"},
+    )
+    context = Context(paths=tmp_path, config=config)
+
+    parent = _create_model("SELECT 1 as id", meta="MODEL (name silver_db.sch.b)")
+    parent = t.cast(SqlModel, context.upsert_model(parent))
+
+    child = _create_model("SELECT id FROM silver_db.sch.b", meta="MODEL (name gold_db.sch.a)")
+    child = t.cast(SqlModel, context.upsert_model(child))
+
+    test_file = tmp_path / "tests" / "test_parameterized_model_names.yaml"
+    test_file.write_text(
+        """
+test_parameterized_model_names:
+  model: {{ var('gold') }}.sch.a
+  inputs:
+    {{ var('silver') }}.sch.b:
+      - id: 1
+  outputs:
+    query:
+      - id: 1
+        """
+    )
+
+    results = context.test()
+
+    assert not results.failures
+    assert not results.errors
+
+    # The example project has one test and we added another one above
+    assert len(results.successes) == 2
+
+
 def test_test_generation(tmp_path: Path) -> None:
     init_example_project(tmp_path, dialect="duckdb")