making it key for model

Author: cmgoffena13

sqlmesh/core/context.py

@@ -40,13 +40,13 @@
 import time
 import traceback
 import typing as t
+from datetime import datetime
 from functools import cached_property
 from io import StringIO
 from itertools import chain
 from pathlib import Path
 from shutil import rmtree
 from types import MappingProxyType
-from datetime import datetime
 
 from sqlglot import Dialect, exp
 from sqlglot.helper import first
@@ -63,6 +63,7 @@
 )
 from sqlmesh.core.config.connection import ConnectionConfig
 from sqlmesh.core.config.loader import C
+from sqlmesh.core.config.model import ModelDefaultsConfig
 from sqlmesh.core.config.root import RegexKeyDict
 from sqlmesh.core.console import get_console
 from sqlmesh.core.context_diff import ContextDiff
@@ -76,24 +77,23 @@
 )
 from sqlmesh.core.engine_adapter import EngineAdapter
 from sqlmesh.core.environment import Environment, EnvironmentNamingInfo, EnvironmentStatements
-from sqlmesh.core.loader import Loader
 from sqlmesh.core.linter.definition import AnnotatedRuleViolation, Linter
 from sqlmesh.core.linter.rules import BUILTIN_RULES
+from sqlmesh.core.loader import Loader
 from sqlmesh.core.macros import ExecutableOrMacro, macro
 from sqlmesh.core.metric import Metric, rewrite
 from sqlmesh.core.model import Model, update_model_schemas
-from sqlmesh.core.config.model import ModelDefaultsConfig
 from sqlmesh.core.notification_target import (
     NotificationEvent,
     NotificationTarget,
     NotificationTargetManager,
 )
-from sqlmesh.core.plan import Plan, PlanBuilder, SnapshotIntervals, PlanExplainer
+from sqlmesh.core.plan import Plan, PlanBuilder, PlanExplainer, SnapshotIntervals
 from sqlmesh.core.plan.definition import UserProvidedFlags
 from sqlmesh.core.reference import ReferenceGraph
-from sqlmesh.core.scheduler import Scheduler, CompletionStatus
+from sqlmesh.core.scheduler import CompletionStatus, Scheduler
 from sqlmesh.core.schema_loader import create_external_models_file
-from sqlmesh.core.selector import Selector, NativeSelector
+from sqlmesh.core.selector import NativeSelector, Selector
 from sqlmesh.core.snapshot import (
     DeployabilityIndex,
     Snapshot,
@@ -111,42 +111,42 @@
 )
 from sqlmesh.core.table_diff import TableDiff
 from sqlmesh.core.test import (
-    ModelTextTestResult,
     ModelTestMetadata,
+    ModelTextTestResult,
     generate_test,
     run_tests,
 )
 from sqlmesh.core.user import User
 from sqlmesh.utils import UniqueKeyDict, Verbosity
 from sqlmesh.utils.concurrency import concurrent_apply_to_values
+from sqlmesh.utils.config import print_config
 from sqlmesh.utils.dag import DAG
 from sqlmesh.utils.date import (
     TimeLike,
-    to_timestamp,
     format_tz_datetime,
-    now_timestamp,
+    make_exclusive,
     now,
+    now_timestamp,
     to_datetime,
-    make_exclusive,
+    to_timestamp,
 )
 from sqlmesh.utils.errors import (
     CircuitBreakerError,
     ConfigError,
+    LinterError,
     PlanError,
     SQLMeshError,
     UncategorizedPlanError,
-    LinterError,
 )
-from sqlmesh.utils.config import print_config
 from sqlmesh.utils.jinja import JinjaMacroRegistry
 
 if t.TYPE_CHECKING:
     import pandas as pd
     from typing_extensions import Literal
 
     from sqlmesh.core.engine_adapter._typing import (
-        BigframeSession,
         DF,
+        BigframeSession,
         PySparkDataFrame,
         PySparkSession,
         SnowparkSession,
@@ -389,6 +389,7 @@ def __init__(
         self._standalone_audits: UniqueKeyDict[str, StandaloneAudit] = UniqueKeyDict(
             "standaloneaudits"
         )
+        self._models_with_tests: t.Set[str] = set()
         self._macros: UniqueKeyDict[str, ExecutableOrMacro] = UniqueKeyDict("macros")
         self._metrics: UniqueKeyDict[str, Metric] = UniqueKeyDict("metrics")
         self._jinja_macros = JinjaMacroRegistry()
@@ -638,6 +639,7 @@ def load(self, update_schemas: bool = True) -> GenericContext[C]:
             self._requirements.update(project.requirements)
             self._excluded_requirements.update(project.excluded_requirements)
             self._environment_statements.extend(project.environment_statements)
+            self._models_with_tests.update(project.models_with_tests)
 
             config = loader.config
             self._linters[config.project] = Linter.from_rules(
@@ -1040,6 +1042,11 @@ def standalone_audits(self) -> MappingProxyType[str, StandaloneAudit]:
         """Returns all registered standalone audits in this context."""
         return MappingProxyType(self._standalone_audits)
 
+    @property
+    def models_with_tests(self) -> t.Set[str]:
+        """Returns all models with tests in this context."""
+        return self._models_with_tests
+
     @property
     def snapshots(self) -> t.Dict[str, Snapshot]:
         """Generates and returns snapshots based on models registered in this context.

sqlmesh/core/linter/rules/builtin.py

@@ -4,7 +4,6 @@
 
 import typing as t
 
-from ruamel.yaml import YAML
 from sqlglot.expressions import Star
 from sqlglot.helper import subclasses
 
@@ -138,25 +137,7 @@ def check_model(self, model: Model) -> t.Optional[RuleViolation]:
         if isinstance(model, ExternalModel):
             return None
 
-        test_dir = self.context.path / "tests"
-        found_test = False
-
-        yaml_parser = YAML(typ="safe")
-        for test_file in test_dir.rglob("*.yaml"):
-            try:
-                test_data = yaml_parser.load(test_file) or {}
-            except Exception:
-                # Skip files with Jinja templating or other parse errors
-                continue
-
-            for _, test_config in test_data.items():
-                if test_config.get("model") == model.name:
-                    found_test = True
-                    break
-            if found_test:
-                break
-
-        if not found_test:
+        if model.name not in self.context.models_with_tests:
             return self.violation(
                 violation_msg=f"Model {model.name} is missing unit test(s). Please add in the tests/ directory."
             )

sqlmesh/core/loader.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import abc
+import concurrent.futures
 import glob
 import itertools
 import linecache
@@ -10,20 +11,19 @@
 from collections import Counter, defaultdict
 from dataclasses import dataclass
 from pathlib import Path
-from pydantic import ValidationError
-import concurrent.futures
 
-from sqlglot.errors import SqlglotError
+from pydantic import ValidationError
 from sqlglot import exp
+from sqlglot.errors import SqlglotError
 from sqlglot.helper import subclasses
 
 from sqlmesh.core import constants as c
 from sqlmesh.core.audit import Audit, ModelAudit, StandaloneAudit, load_multiple_audits
 from sqlmesh.core.console import Console
 from sqlmesh.core.dialect import parse
 from sqlmesh.core.environment import EnvironmentStatements
-from sqlmesh.core.linter.rule import Rule
 from sqlmesh.core.linter.definition import RuleSet
+from sqlmesh.core.linter.rule import Rule
 from sqlmesh.core.macros import MacroRegistry, macro
 from sqlmesh.core.metric import Metric, MetricMeta, expand_metrics, load_metric_ddl
 from sqlmesh.core.model import (
@@ -40,10 +40,10 @@
 from sqlmesh.utils.errors import ConfigError
 from sqlmesh.utils.jinja import JinjaMacroRegistry, MacroExtractor
 from sqlmesh.utils.metaprogramming import import_python_file
-from sqlmesh.utils.pydantic import validation_error_message
 from sqlmesh.utils.process import create_process_pool_executor
-from sqlmesh.utils.yaml import YAML, load as yaml_load
-
+from sqlmesh.utils.pydantic import validation_error_message
+from sqlmesh.utils.yaml import YAML
+from sqlmesh.utils.yaml import load as yaml_load
 
 if t.TYPE_CHECKING:
     from sqlmesh.core.context import GenericContext
@@ -64,6 +64,8 @@ class LoadedProject:
     excluded_requirements: t.Set[str]
     environment_statements: t.List[EnvironmentStatements]
     user_rules: RuleSet
+    model_test_metadata: t.List[ModelTestMetadata]
+    models_with_tests: t.Set[str]
 
 
 class CacheBase(abc.ABC):
@@ -243,6 +245,14 @@ def load(self) -> LoadedProject:
 
             user_rules = self._load_linting_rules()
 
+            model_test_metadata = self.load_model_tests()
+
+            models_with_tests = {
+                model_test_metadata.model_name for model_test_metadata in model_test_metadata
+            }
+
+            self._models_with_tests = models_with_tests
+
             project = LoadedProject(
                 macros=macros,
                 jinja_macros=jinja_macros,
@@ -254,6 +264,8 @@ def load(self) -> LoadedProject:
                 excluded_requirements=excluded_requirements,
                 environment_statements=environment_statements,
                 user_rules=user_rules,
+                model_test_metadata=model_test_metadata,
+                models_with_tests=models_with_tests,
             )
             return project
 

sqlmesh/core/test/discovery.py

@@ -20,6 +20,10 @@ class ModelTestMetadata(PydanticModel):
     def fully_qualified_test_name(self) -> str:
         return f"{self.path}::{self.test_name}"
 
+    @property
+    def model_name(self) -> str:
+        return self.body["model"]
+
     def __hash__(self) -> int:
         return self.fully_qualified_test_name.__hash__()