[POC] Unit Testing

core/dbt/cli/main.py

@@ -39,6 +39,7 @@
 from dbt.task.show import ShowTask
 from dbt.task.snapshot import SnapshotTask
 from dbt.task.test import TestTask
+from dbt.task.unit_test import UnitTestTask
 
 
 @dataclass
@@ -845,6 +846,52 @@
     return results, success
 
 
+# dbt test
+@cli.command("unit-test")
+@click.pass_context
+@p.defer
+@p.deprecated_defer
+@p.exclude
+@p.fail_fast
+@p.favor_state
+@p.deprecated_favor_state
+@p.indirect_selection
+@p.show_output_format
+@p.profile
+@p.profiles_dir
+@p.project_dir
+@p.select
+@p.selector
+@p.state
+@p.defer_state
+@p.deprecated_state
+@p.store_failures
+@p.target
+@p.target_path
+@p.threads
+@p.vars
+@p.version_check
+@requires.postflight
+@requires.preflight
+@requires.profile
+@requires.project
+@requires.runtime_config
+@requires.manifest
+@requires.unit_test_collection
+def unit_test(ctx, **kwargs):
+    """Runs tests on data in deployed models. Run this after `dbt run`"""
+    task = UnitTestTask(
+        ctx.obj["flags"],
+        ctx.obj["runtime_config"],
+        ctx.obj["manifest"],
+        ctx.obj["unit_test_collection"],
+    )
+
+    results = task.run()
+    success = task.interpret_results(results)
+    return results, success
+
+
 # Support running as a module
 if __name__ == "__main__":
     cli()

core/dbt/cli/requires.py

@@ -23,6 +23,7 @@
 from dbt.profiler import profiler
 from dbt.tracking import active_user, initialize_from_flags, track_run
 from dbt.utils import cast_dict_to_dict_of_strings
+from dbt.parser.unit_tests import UnitTestManifestLoader
 from dbt.plugins import set_up_plugin_manager, get_plugin_manager
 
 from click import Context
@@ -265,3 +266,25 @@
     if len(args0) == 0:
         return outer_wrapper
     return outer_wrapper(args0[0])
+
+
+def unit_test_collection(func):
+    """A decorator used by click command functions for generating a unit test collection provided a manifest"""
+
+    def wrapper(*args, **kwargs):
+        ctx = args[0]
+        assert isinstance(ctx, Context)
+
+        req_strs = ["manifest", "runtime_config"]
+        reqs = [ctx.obj.get(req_str) for req_str in req_strs]
+
+        if None in reqs:
+            raise DbtProjectError("manifest and runtime_config required for unit_test_collection")
+
+        collection = UnitTestManifestLoader.load(ctx.obj["manifest"], ctx.obj["runtime_config"])
+
+        ctx.obj["unit_test_collection"] = collection
+
+        return func(*args, **kwargs)
+
+    return update_wrapper(wrapper, func)

core/dbt/clients/jinja.py

@@ -330,6 +330,26 @@
             return self.call_macro(*args, **kwargs)
 
 
+class UnitTestMacroGenerator(MacroGenerator):
+    # this makes UnitTestMacroGenerator objects callable like functions
+    def __init__(
+        self,
+        macro_generator: MacroGenerator,
+        call_return_value: Any,
+    ) -> None:
+        super().__init__(
+            macro_generator.macro,
+            macro_generator.context,
+            macro_generator.node,
+            macro_generator.stack,
+        )
+        self.call_return_value = call_return_value
+
+    def __call__(self, *args, **kwargs):
+        with self.track_call():
+            return self.call_return_value
+
+
 class QueryStringGenerator(BaseMacroGenerator):
     def __init__(self, template_str: str, context: Dict[str, Any]) -> None:
         super().__init__(context)

core/dbt/compilation.py

@@ -13,7 +13,10 @@
 from dbt.adapters.factory import get_adapter
 from dbt.clients import jinja
 from dbt.clients.system import make_directory
-from dbt.context.providers import generate_runtime_model_context
+from dbt.context.providers import (
+    generate_runtime_model_context,
+    generate_runtime_unit_test_context,
+)
 from dbt.contracts.graph.manifest import Manifest, UniqueID
 from dbt.contracts.graph.nodes import (
     ManifestNode,
@@ -22,6 +25,7 @@
     GraphMemberNode,
     InjectedCTE,
     SeedNode,
+    UnitTestNode,
 )
 from dbt.exceptions import (
     GraphDependencyNotFoundError,
@@ -44,6 +48,7 @@
     names = {
         NodeType.Model: "model",
         NodeType.Test: "test",
+        NodeType.Unit: "unit test",
         NodeType.Snapshot: "snapshot",
         NodeType.Analysis: "analysis",
         NodeType.Macro: "macro",
@@ -289,8 +294,10 @@
         manifest: Manifest,
         extra_context: Dict[str, Any],
     ) -> Dict[str, Any]:
-
-        context = generate_runtime_model_context(node, self.config, manifest)
+        if isinstance(node, UnitTestNode):
+            context = generate_runtime_unit_test_context(node, self.config, manifest)
+        else:
+            context = generate_runtime_model_context(node, self.config, manifest)
         context.update(extra_context)
 
         if isinstance(node, GenericTestNode):

core/dbt/context/providers.py

@@ -1,4 +1,5 @@
 import abc
+from copy import deepcopy
 import os
 from typing import (
     Callable,
@@ -17,7 +18,7 @@
 from dbt.adapters.base.column import Column
 from dbt.adapters.factory import get_adapter, get_adapter_package_names, get_adapter_type_names
 from dbt.clients import agate_helper
-from dbt.clients.jinja import get_rendered, MacroGenerator, MacroStack
+from dbt.clients.jinja import get_rendered, MacroGenerator, MacroStack, UnitTestMacroGenerator
 from dbt.config import RuntimeConfig, Project
 from dbt.constants import SECRET_ENV_PREFIX, DEFAULT_ENV_PLACEHOLDER
 from dbt.context.base import contextmember, contextproperty, Var
@@ -39,6 +40,7 @@
     RefArgs,
     AccessType,
     SemanticModel,
+    UnitTestNode,
 )
 from dbt.contracts.graph.metrics import MetricReference, ResolvedMetricReference
 from dbt.contracts.graph.unparsed import NodeVersion
@@ -566,6 +568,17 @@
             return super().create_relation(target_model)
 
 
+class RuntimeUnitTestRefResolver(RuntimeRefResolver):
+    def resolve(
+        self,
+        target_name: str,
+        target_package: Optional[str] = None,
+        target_version: Optional[NodeVersion] = None,
+    ) -> RelationProxy:
+        target_name = f"{self.model.name}__{target_name}"
+        return super().resolve(target_name, target_package, target_version)
+
+
 # `source` implementations
 class ParseSourceResolver(BaseSourceResolver):
     def resolve(self, source_name: str, table_name: str):
@@ -670,6 +683,22 @@
     pass
 
 
+class UnitTestVar(RuntimeVar):
+    def __init__(
+        self,
+        context: Dict[str, Any],
+        config: RuntimeConfig,
+        node: Resource,
+    ) -> None:
+        config_copy = None
+        assert isinstance(node, UnitTestNode)
+        if node.overrides and node.overrides.vars:
+            config_copy = deepcopy(config)
+            config_copy.cli_vars.update(node.overrides.vars)
+
+        super().__init__(context, config_copy or config, node=node)
+
+
 # Providers
 class Provider(Protocol):
     execute: bool
@@ -711,6 +740,16 @@
     metric = RuntimeMetricResolver
 
 
+class RuntimeUnitTestProvider(Provider):
+    execute = True
+    Config = RuntimeConfigObject
+    DatabaseWrapper = RuntimeDatabaseWrapper
+    Var = UnitTestVar
+    ref = RuntimeUnitTestRefResolver
+    source = RuntimeSourceResolver  # TODO: RuntimeUnitTestSourceResolver
+    metric = RuntimeMetricResolver
+
+
 class OperationProvider(RuntimeProvider):
     ref = OperationRefResolver
 
@@ -1359,7 +1398,7 @@
 
     @contextproperty
     def pre_hooks(self) -> List[Dict[str, Any]]:
-        if self.model.resource_type in [NodeType.Source, NodeType.Test]:
+        if self.model.resource_type in [NodeType.Source, NodeType.Test, NodeType.Unit]:
             return []
         # TODO CT-211
         return [
@@ -1368,7 +1407,7 @@
 
     @contextproperty
     def post_hooks(self) -> List[Dict[str, Any]]:
-        if self.model.resource_type in [NodeType.Source, NodeType.Test]:
+        if self.model.resource_type in [NodeType.Source, NodeType.Test, NodeType.Unit]:
             return []
         # TODO CT-211
         return [
@@ -1461,6 +1500,25 @@
             return None
 
 
+class UnitTestContext(ModelContext):
+    model: UnitTestNode
+
+    @contextmember
+    def env_var(self, var: str, default: Optional[str] = None) -> str:
+        """The env_var() function. Return the overriden unit test environment variable named 'var'.
+
+        If there is no unit test override, return the environment variable named 'var'.
+
+        If there is no such environment variable set, return the default.
+
+        If the default is None, raise an exception for an undefined variable.
+        """
+        if self.model.overrides and var in self.model.overrides.env_vars:
+            return self.model.overrides.env_vars[var]
+        else:
+            return super().env_var(var, default)
+
+
 # This is called by '_context_for', used in 'render_with_context'
 def generate_parser_model_context(
     model: ManifestNode,
@@ -1505,6 +1563,24 @@
     return ctx.to_dict()
 
 
+def generate_runtime_unit_test_context(
+    unit_test: UnitTestNode,
+    config: RuntimeConfig,
+    manifest: Manifest,
+) -> Dict[str, Any]:
+    ctx = UnitTestContext(unit_test, config, manifest, RuntimeUnitTestProvider(), None)
+    ctx_dict = ctx.to_dict()
+
+    if unit_test.overrides and unit_test.overrides.macros:
+        for macro_name, macro_value in unit_test.overrides.macros.items():
+            context_value = ctx_dict.get(macro_name)
+            if isinstance(context_value, MacroGenerator):
+                ctx_dict[macro_name] = UnitTestMacroGenerator(context_value, macro_value)
+            else:
+                ctx_dict[macro_name] = macro_value
+    return ctx_dict
+
+
 class ExposureRefResolver(BaseResolver):
     def __call__(self, *args, **kwargs) -> str:
         package = None

core/dbt/contracts/graph/model_config.py

@@ -661,6 +661,7 @@ def finalize_and_validate(self):
     NodeType.Source: SourceConfig,
     NodeType.Seed: SeedConfig,
     NodeType.Test: TestConfig,
+    NodeType.Unit: TestConfig,
     NodeType.Model: NodeConfig,
     NodeType.Snapshot: SnapshotConfig,
 }

core/dbt/contracts/graph/nodes.py

@@ -34,6 +34,7 @@
     UnparsedSourceDefinition,
     UnparsedSourceTableDefinition,
     UnparsedColumn,
+    UnparsedUnitTestOverrides,
 )
 from dbt.contracts.graph.node_args import ModelNodeArgs
 from dbt.contracts.util import Replaceable, AdditionalPropertiesMixin
@@ -942,13 +943,27 @@
     namespace: Optional[str] = None
 
 
+@dataclass
+class UnitTestMetadata(dbtClassMixin, Replaceable):
+    # kwargs are the args that are left in the test builder after
+    # removing configs. They are set from the test builder when
+    # the test node is created.
+    kwargs: Dict[str, Any] = field(default_factory=dict)
+    namespace: Optional[str] = None
+
+
 # This has to be separated out because it has no default and so
 # has to be included as a superclass, not an attribute
 @dataclass
 class HasTestMetadata(dbtClassMixin):
     test_metadata: TestMetadata
 
 
+@dataclass
+class HasUnitTestMetadata(dbtClassMixin):
+    unit_test_metadata: UnitTestMetadata
+
+
 @dataclass
 class GenericTestNode(TestShouldStoreFailures, CompiledNode, HasTestMetadata):
     resource_type: NodeType = field(metadata={"restrict": [NodeType.Test]})
@@ -970,6 +985,17 @@
         return "generic"
 
 
+@dataclass
+class UnitTestNode(CompiledNode):
+    resource_type: NodeType = field(metadata={"restrict": [NodeType.Unit]})
+    attached_node: Optional[str] = None
+    overrides: Optional[UnparsedUnitTestOverrides] = None
+
+    @property
+    def test_node_type(self):
+        return "unit"
+
+
 # ====================================
 # Snapshot node
 # ====================================
@@ -1628,6 +1654,7 @@
     SqlNode,
     GenericTestNode,
     SnapshotNode,
+    UnitTestNode,
 ]
 
 # All SQL nodes plus SeedNode (csv files)
@@ -1657,7 +1684,4 @@
     Group,
 ]
 
-TestNode = Union[
-    SingularTestNode,
-    GenericTestNode,
-]
+TestNode = Union[SingularTestNode, GenericTestNode]