diff --git a/src/taskgraph/morph.py b/src/taskgraph/morph.py index a28250c37..c48831778 100644 --- a/src/taskgraph/morph.py +++ b/src/taskgraph/morph.py @@ -33,6 +33,12 @@ logger = logging.getLogger(__name__) MAX_ROUTES = 10 +registered_morphs = [] + + +def register_morph(func): + registered_morphs.append(func) + def amend_taskgraph(taskgraph, label_to_taskid, to_add): """Add the given tasks to the taskgraph, returning a new taskgraph""" @@ -156,6 +162,7 @@ def make_index_task(parent_task, taskgraph, label_to_taskid, parameters, graph_c return task, taskgraph, label_to_taskid +@register_morph def add_index_tasks(taskgraph, label_to_taskid, parameters, graph_config): """ The TaskCluster queue only allows 10 routes on a task, but we have tasks @@ -196,8 +203,9 @@ def _get_morph_url(): return f"{taskgraph_repo}/raw-file/{taskgraph_rev}/src/taskgraph/morph.py" +@register_morph def add_code_review_task(taskgraph, label_to_taskid, parameters, graph_config): - logger.debug("Morphing: adding index tasks") + logger.debug("Morphing: adding code review task") review_config = parameters.get("code-review") if not review_config: @@ -256,12 +264,7 @@ def add_code_review_task(taskgraph, label_to_taskid, parameters, graph_config): def morph(taskgraph, label_to_taskid, parameters, graph_config): """Apply all morphs""" - morphs = [ - add_index_tasks, - add_code_review_task, - ] - - for m in morphs: + for m in registered_morphs: taskgraph, label_to_taskid = m( taskgraph, label_to_taskid, parameters, graph_config ) diff --git a/test/test_morph.py b/test/test_morph.py index 0837bf8d6..56ef9a4aa 100644 --- a/test/test_morph.py +++ b/test/test_morph.py @@ -72,3 +72,18 @@ def test_make_index_tasks(make_taskgraph, graph_config): # check the scope summary assert index_task.task["scopes"] == ["index:insert-task:gecko.v2.mozilla-central.*"] + + +def test_register_morph(monkeypatch, make_taskgraph): + taskgraph, label_to_taskid = make_taskgraph({}) + monkeypatch.setattr(morph, "registered_morphs", []) + + @morph.register_morph + def fake_morph(taskgraph, label_to_taskid, *args): + label_to_taskid.setdefault("count", 0) + label_to_taskid["count"] += 1 + return taskgraph, label_to_taskid + + assert label_to_taskid == {} + morph.morph(taskgraph, label_to_taskid, None, None) + assert label_to_taskid == {"count": 1}