diff --git a/experimental/python/databricks/bundles/build.py b/experimental/python/databricks/bundles/build.py index 949ac39ae..6a867ef62 100644 --- a/experimental/python/databricks/bundles/build.py +++ b/experimental/python/databricks/bundles/build.py @@ -128,6 +128,7 @@ def _apply_mutators( ("resources", "jobs", resource_name), location ) resources.jobs[resource_name] = new_job + job = new_job except Exception as exc: mutator_name = mutator.function.__name__ diff --git a/experimental/python/databricks_tests/test_build.py b/experimental/python/databricks_tests/test_build.py index 427d447b6..2d08e138a 100644 --- a/experimental/python/databricks_tests/test_build.py +++ b/experimental/python/databricks_tests/test_build.py @@ -1,9 +1,11 @@ import sys +from dataclasses import replace from io import StringIO from pathlib import Path from databricks.bundles.build import ( _append_resources, + _apply_mutators, _Args, _Conf, _load_object, @@ -21,6 +23,7 @@ from databricks.bundles.core import ( Location, Resources, Severity, + job_mutator, ) from databricks.bundles.jobs import Job @@ -299,6 +302,69 @@ def test_conf_from_dict(): ) +def test_mutators(): + bundle = Bundle(target="default") + resources = Resources() + resources.add_job("job_0", Job(tags={"tag": "value"})) + + @job_mutator + def add_first_tag(bundle: Bundle, job: Job) -> Job: + tags = bundle.resolve_variable(job.tags) + + return replace(job, tags={"first": "tag", **tags}) + + @job_mutator + def add_second_tag(bundle: Bundle, job: Job) -> Job: + tags = bundle.resolve_variable(job.tags) + + return replace(job, tags={"second": "tag", **tags}) + + new_resources, diagnostics = _apply_mutators( + bundle=bundle, + resources=resources, + mutator_functions=[add_first_tag, add_second_tag], + ) + + # add_second_tag is the last mutator that has modified a job + expected_location = Location.from_callable(add_second_tag.function) + + assert not diagnostics.has_error() + assert new_resources._locations[("resources", "jobs", "job_0")] == expected_location + assert new_resources.jobs["job_0"].tags == { + "first": "tag", + "second": "tag", + "tag": "value", + } + + +def test_mutators_unmodified(): + bundle = Bundle(target="default") + + resources = Resources() + resources.add_job("job_0", Job(description="description")) + + @job_mutator + def mutator_1(job: Job) -> Job: + return replace(job, description="updated description") + + @job_mutator + def mutator_2(job: Job) -> Job: + return job + + new_resources, diagnostics = _apply_mutators( + bundle=bundle, + resources=resources, + mutator_functions=[mutator_1, mutator_2], + ) + + # despite mutator_2 being called last, it doesn't change the job, and we should use location of mutator_1 + expected_location = Location.from_callable(mutator_1.function) + + assert not diagnostics.has_error() + assert new_resources._locations[("resources", "jobs", "job_0")] == expected_location + assert new_resources.jobs["job_0"].description == "updated description" + + def test_load_resources(): bundle = Bundle(target="default")