From 6725490052846517af4802f56492da597900e836 Mon Sep 17 00:00:00 2001 From: Ben Cassell Date: Mon, 4 Dec 2023 16:23:36 -0800 Subject: [PATCH] fix linting and hopefully tests --- .python-version | 1 + tests/unit/macros/base.py | 88 +++++++++++++++++-- .../macros/relations/test_table_macros.py | 3 - tests/unit/macros/test_adapters_macros.py | 2 +- tests/unit/macros/test_python_macros.py | 10 ++- 5 files changed, 89 insertions(+), 15 deletions(-) create mode 100644 .python-version diff --git a/.python-version b/.python-version new file mode 100644 index 000000000..2c0733315 --- /dev/null +++ b/.python-version @@ -0,0 +1 @@ +3.11 diff --git a/tests/unit/macros/base.py b/tests/unit/macros/base.py index 2df25c642..3d59d22bd 100644 --- a/tests/unit/macros/base.py +++ b/tests/unit/macros/base.py @@ -1,4 +1,5 @@ import re +from typing import Any, Dict from mock import Mock import pytest from jinja2 import Environment, FileSystemLoader, PackageLoader, Template @@ -15,18 +16,27 @@ def __init__(self, template, context, relation): class MacroTestBase: @pytest.fixture(autouse=True) def config(self, context) -> dict: - local_config = {} + """ + Anything you put in this dict will be returned by config in the rendered template + """ + local_config: Dict[str, Any] = {} context["config"].get = lambda key, default=None, **kwargs: local_config.get(key, default) return local_config @pytest.fixture(autouse=True) def var(self, context) -> dict: - local_var = {} + """ + Anything you put in this dict will be returned by config in the rendered template + """ + local_var: Dict[str, Any] = {} context["var"] = lambda key, default=None, **kwargs: local_var.get(key, default) return local_var @pytest.fixture(scope="class") def default_context(self) -> dict: + """ + This is the default context used in all tests. + """ context = { "validation": Mock(), "model": Mock(), @@ -36,12 +46,16 @@ def default_context(self) -> dict: "adapter": Mock(), "var": Mock(), "return": lambda r: r, + "is_incremental": Mock(return_value=False), } return context @pytest.fixture(scope="class") def spark_env(self) -> Environment: + """ + The environment used for rendering dbt-spark macros + """ return Environment( loader=PackageLoader("dbt.include.spark", "macros"), extensions=["jinja2.ext.do"], @@ -49,18 +63,33 @@ def spark_env(self) -> Environment: @pytest.fixture(scope="class") def spark_template_names(self) -> list: + """ + The list of Spark templates to load for the tests. + Use this if your macro relies on macros defined in templates we inherit from dbt-spark. + """ return ["adapters.sql"] @pytest.fixture(scope="class") def spark_context(self, default_context, spark_env, spark_template_names) -> dict: + """ + Adds all the requested Spark macros to the context + """ return self.build_up_context(default_context, spark_env, spark_template_names) @pytest.fixture(scope="class") def macro_folders_to_load(self) -> list: + """ + This is a list of folders from which we look to load Databricks macro templates. + All folders are relative to the dbt/include/databricks folder. + Folders will be searched for in the order they are listed here, in case of name collisions. + """ return ["macros"] @pytest.fixture(scope="class") def databricks_env(self, macro_folders_to_load) -> Environment: + """ + The environment used for rendering Databricks macros + """ return Environment( loader=FileSystemLoader( [f"dbt/include/databricks/{folder}" for folder in macro_folders_to_load] @@ -70,15 +99,28 @@ def databricks_env(self, macro_folders_to_load) -> Environment: @pytest.fixture(scope="class") def databricks_template_names(self) -> list: + """ + The list of databricks templates to load for referencing imported macros in the + tests. Do not include the template you specify in template_name. Use this when you need a + macro defined in a template other than the one you render for the test. + + Ex: If you are testing the python.sql template, you will also need to load ["adapters.sql"] + """ return [] @pytest.fixture(scope="class") def databricks_context(self, spark_context, databricks_env, databricks_template_names) -> dict: + """ + Adds all the requested Databricks macros to the context + """ if not databricks_template_names: return spark_context return self.build_up_context(spark_context, databricks_env, databricks_template_names) def build_up_context(self, context, env, template_names): + """ + Adds macros from the supplied env and template names to the context. + """ new_context = context.copy() for template_name in template_names: template = env.get_template(template_name, globals=context) @@ -86,16 +128,22 @@ def build_up_context(self, context, env, template_names): return new_context - @pytest.fixture - def context(self, databricks_context) -> dict: - return databricks_context.copy() - @pytest.fixture(scope="class") def template_name(self) -> str: + """ + The name of the Databricks template you want to test, not including the path. + + Example: "adapters.sql" + """ raise NotImplementedError("Must be implemented by subclasses") @pytest.fixture - def template(self, template_name, context, databricks_env) -> Template: + def template(self, template_name, databricks_context, databricks_env) -> Template: + """ + This creates the template you will test against. + You generally don't want to override this. + """ + context = databricks_context.copy() current_template = databricks_env.get_template(template_name, globals=context) def dispatch(macro_name, macro_namespace=None, packages=None): @@ -110,8 +158,21 @@ def dispatch(macro_name, macro_namespace=None, packages=None): return current_template + @pytest.fixture + def context(self, template) -> dict: + """ + Access to the context used to render the template. + Modification of the context will work for mocking adapter calls, but may not work for + mocking macros. + If you need to mock a macro, see the use of is_incremental in default_context. + """ + return template.globals + @pytest.fixture(scope="class") def relation(self): + """ + Dummy relation to use in tests. + """ data = { "path": { "database": "some_database", @@ -125,15 +186,28 @@ def relation(self): @pytest.fixture def template_bundle(self, template, context, relation): + """ + Bundles up the compiled template, its context, and a dummy relation. + """ context["model"].alias = relation.identifier return TemplateBundle(template, context, relation) def run_macro_raw(self, template, name, *args): + """ + Run the named macro from a template, and return the rendered value. + """ return getattr(template.module, name)(*args) def run_macro(self, template, name, *args): + """ + Run the named macro from a template, and return the rendered value. + This version strips off extra whitespace and newlines. + """ value = self.run_macro_raw(template, name, *args) return re.sub(r"\s\s+", " ", value).strip() def render_bundle(self, template_bundle, name, *args): + """ + Convenience method for macros that take a relation as a first argument. + """ return self.run_macro(template_bundle.template, name, template_bundle.relation, *args) diff --git a/tests/unit/macros/relations/test_table_macros.py b/tests/unit/macros/relations/test_table_macros.py index b4ddef893..5e91a3b35 100644 --- a/tests/unit/macros/relations/test_table_macros.py +++ b/tests/unit/macros/relations/test_table_macros.py @@ -1,6 +1,3 @@ -from mock import Mock -from jinja2 import Environment, FileSystemLoader, PackageLoader -import re import pytest from tests.unit.macros.base import MacroTestBase diff --git a/tests/unit/macros/test_adapters_macros.py b/tests/unit/macros/test_adapters_macros.py index bce535f0f..e3850cd9e 100644 --- a/tests/unit/macros/test_adapters_macros.py +++ b/tests/unit/macros/test_adapters_macros.py @@ -200,7 +200,7 @@ def render_constraint_sql(self, template_bundle, constraint, *args): "get_constraint_sql", template_bundle.relation, constraint, - *args + *args, ) @pytest.fixture(scope="class") diff --git a/tests/unit/macros/test_python_macros.py b/tests/unit/macros/test_python_macros.py index 246b520c0..b59b2a4e6 100644 --- a/tests/unit/macros/test_python_macros.py +++ b/tests/unit/macros/test_python_macros.py @@ -1,3 +1,4 @@ +from jinja2 import Template from mock import MagicMock from tests.unit.macros.base import MacroTestBase @@ -6,7 +7,7 @@ class TestPythonMacros(MacroTestBase): @pytest.fixture(scope="class", autouse=True) - def modify_context(self, default_context) -> dict: + def modify_context(self, default_context) -> None: default_context["model"] = MagicMock() d = {"alias": "schema"} default_context["model"].__getitem__.side_effect = d.__getitem__ @@ -32,15 +33,16 @@ def test_py_get_writer__specified_file_format(self, config, template): def test_py_get_writer__specified_location_root(self, config, template, context): config["location_root"] = "s3://fake_location" - context["is_incremental"] = MagicMock(return_value=False) result = self.run_macro_raw(template, "py_get_writer_options") expected = '.format("delta")\n.option("path", "s3://fake_location/schema")' assert result == expected - def test_py_get_writer__specified_location_root_on_incremental(self, config, template, context): + def test_py_get_writer__specified_location_root_on_incremental( + self, config, template: Template, context + ): config["location_root"] = "s3://fake_location" - context["is_incremental"] = MagicMock(return_value=True) + context["is_incremental"].return_value = True result = self.run_macro_raw(template, "py_get_writer_options") expected = '.format("delta")\n.option("path", "s3://fake_location/schema__dbt_tmp")'