[Python] Fix issues with multiple mutators

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__
 

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 @@
     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")