sqlmesh.core.test.context
1from __future__ import annotations 2 3import typing as t 4from functools import cached_property 5 6from sqlmesh import Model 7from sqlmesh.core.context import ExecutionContext 8from sqlmesh.core.engine_adapter import EngineAdapter 9from sqlmesh.core.test.definition import ModelTest 10from sqlmesh.utils import UniqueKeyDict 11 12 13class TestExecutionContext(ExecutionContext): 14 """The context needed to execute a Python model test. 15 16 Args: 17 engine_adapter: The engine adapter to execute queries against. 18 models: All upstream models to use for expansion and mapping of physical locations. 19 """ 20 21 __test__ = False # prevent pytest trying to collect this as a test class 22 23 def __init__( 24 self, 25 engine_adapter: EngineAdapter, 26 models: UniqueKeyDict[str, Model], 27 test: ModelTest, 28 default_dialect: t.Optional[str] = None, 29 default_catalog: t.Optional[str] = None, 30 variables: t.Optional[t.Dict[str, t.Any]] = None, 31 blueprint_variables: t.Optional[t.Dict[str, t.Any]] = None, 32 ): 33 self._engine_adapter = engine_adapter 34 self._models = models 35 self._test = test 36 self._default_catalog = default_catalog 37 self._default_dialect = default_dialect 38 self._variables = variables or {} 39 self._blueprint_variables = variables or {} 40 41 @cached_property 42 def _model_tables(self) -> t.Dict[str, str]: 43 """Returns a mapping of model names to tables.""" 44 45 # Include upstream dependencies to ensure they can be resolved during test execution 46 return { 47 name: self._test._test_fixture_table(name).sql() 48 for normalized_model_name, model in self._models.items() 49 for name in [normalized_model_name, *model.depends_on] 50 } 51 52 def with_variables( 53 self, 54 variables: t.Dict[str, t.Any], 55 blueprint_variables: t.Optional[t.Dict[str, t.Any]] = None, 56 ) -> TestExecutionContext: 57 """Returns a new TestExecutionContext with additional variables.""" 58 return TestExecutionContext( 59 self._engine_adapter, 60 self._models, 61 self._test, 62 self._default_dialect, 63 self._default_catalog, 64 variables=variables, 65 blueprint_variables=blueprint_variables, 66 )
14class TestExecutionContext(ExecutionContext): 15 """The context needed to execute a Python model test. 16 17 Args: 18 engine_adapter: The engine adapter to execute queries against. 19 models: All upstream models to use for expansion and mapping of physical locations. 20 """ 21 22 __test__ = False # prevent pytest trying to collect this as a test class 23 24 def __init__( 25 self, 26 engine_adapter: EngineAdapter, 27 models: UniqueKeyDict[str, Model], 28 test: ModelTest, 29 default_dialect: t.Optional[str] = None, 30 default_catalog: t.Optional[str] = None, 31 variables: t.Optional[t.Dict[str, t.Any]] = None, 32 blueprint_variables: t.Optional[t.Dict[str, t.Any]] = None, 33 ): 34 self._engine_adapter = engine_adapter 35 self._models = models 36 self._test = test 37 self._default_catalog = default_catalog 38 self._default_dialect = default_dialect 39 self._variables = variables or {} 40 self._blueprint_variables = variables or {} 41 42 @cached_property 43 def _model_tables(self) -> t.Dict[str, str]: 44 """Returns a mapping of model names to tables.""" 45 46 # Include upstream dependencies to ensure they can be resolved during test execution 47 return { 48 name: self._test._test_fixture_table(name).sql() 49 for normalized_model_name, model in self._models.items() 50 for name in [normalized_model_name, *model.depends_on] 51 } 52 53 def with_variables( 54 self, 55 variables: t.Dict[str, t.Any], 56 blueprint_variables: t.Optional[t.Dict[str, t.Any]] = None, 57 ) -> TestExecutionContext: 58 """Returns a new TestExecutionContext with additional variables.""" 59 return TestExecutionContext( 60 self._engine_adapter, 61 self._models, 62 self._test, 63 self._default_dialect, 64 self._default_catalog, 65 variables=variables, 66 blueprint_variables=blueprint_variables, 67 )
The context needed to execute a Python model test.
Arguments:
- engine_adapter: The engine adapter to execute queries against.
- models: All upstream models to use for expansion and mapping of physical locations.
TestExecutionContext( engine_adapter: sqlmesh.core.engine_adapter.base.EngineAdapter, models: sqlmesh.utils.UniqueKeyDict[str, typing.Union[sqlmesh.core.model.definition.SqlModel, sqlmesh.core.model.definition.SeedModel, sqlmesh.core.model.definition.PythonModel, sqlmesh.core.model.definition.ExternalModel]], test: sqlmesh.core.test.definition.ModelTest, default_dialect: Optional[str] = None, default_catalog: Optional[str] = None, variables: Optional[Dict[str, Any]] = None, blueprint_variables: Optional[Dict[str, Any]] = None)
24 def __init__( 25 self, 26 engine_adapter: EngineAdapter, 27 models: UniqueKeyDict[str, Model], 28 test: ModelTest, 29 default_dialect: t.Optional[str] = None, 30 default_catalog: t.Optional[str] = None, 31 variables: t.Optional[t.Dict[str, t.Any]] = None, 32 blueprint_variables: t.Optional[t.Dict[str, t.Any]] = None, 33 ): 34 self._engine_adapter = engine_adapter 35 self._models = models 36 self._test = test 37 self._default_catalog = default_catalog 38 self._default_dialect = default_dialect 39 self._variables = variables or {} 40 self._blueprint_variables = variables or {}
def
with_variables( self, variables: Dict[str, Any], blueprint_variables: Optional[Dict[str, Any]] = None) -> TestExecutionContext:
53 def with_variables( 54 self, 55 variables: t.Dict[str, t.Any], 56 blueprint_variables: t.Optional[t.Dict[str, t.Any]] = None, 57 ) -> TestExecutionContext: 58 """Returns a new TestExecutionContext with additional variables.""" 59 return TestExecutionContext( 60 self._engine_adapter, 61 self._models, 62 self._test, 63 self._default_dialect, 64 self._default_catalog, 65 variables=variables, 66 blueprint_variables=blueprint_variables, 67 )
Returns a new TestExecutionContext with additional variables.