1# This Source Code Form is subject to the terms of the Mozilla Public
2# License, v. 2.0. If a copy of the MPL was not distributed with this
3# file, You can obtain one at http://mozilla.org/MPL/2.0/.
4
5from __future__ import absolute_import, print_function, unicode_literals
6
7import pytest
8import six
9import unittest
10from mozunit import main
11
12from taskgraph.generator import TaskGraphGenerator, Kind, load_tasks_for_kind
13from taskgraph.optimize import OptimizationStrategy
14from taskgraph.config import GraphConfig
15from taskgraph.util.templates import merge
16from taskgraph import (
17    generator,
18    graph,
19    optimize as optimize_mod,
20    target_tasks as target_tasks_mod,
21)
22
23
24def fake_loader(kind, path, config, parameters, loaded_tasks):
25    for i in range(3):
26        dependencies = {}
27        if i >= 1:
28            dependencies["prev"] = "{}-t-{}".format(kind, i - 1)
29
30        task = {
31            "kind": kind,
32            "label": "{}-t-{}".format(kind, i),
33            "description": "{} task {}".format(kind, i),
34            "attributes": {"_tasknum": six.text_type(i)},
35            "task": {
36                "i": i,
37                "metadata": {"name": "t-{}".format(i)},
38                "deadline": "soon",
39            },
40            "dependencies": dependencies,
41        }
42        if "job-defaults" in config:
43            task = merge(config["job-defaults"], task)
44        yield task
45
46
47class FakeKind(Kind):
48    def _get_loader(self):
49        return fake_loader
50
51    def load_tasks(self, parameters, loaded_tasks, write_artifacts):
52        FakeKind.loaded_kinds.append(self.name)
53        return super(FakeKind, self).load_tasks(
54            parameters, loaded_tasks, write_artifacts
55        )
56
57
58class WithFakeKind(TaskGraphGenerator):
59    def _load_kinds(self, graph_config, target_kind=None):
60        for kind_name, cfg in self.parameters["_kinds"]:
61            config = {
62                "transforms": [],
63            }
64            if cfg:
65                config.update(cfg)
66            yield FakeKind(kind_name, "/fake", config, graph_config)
67
68
69def fake_load_graph_config(root_dir):
70    graph_config = GraphConfig(
71        {"trust-domain": "test-domain", "taskgraph": {}}, root_dir
72    )
73    graph_config.__dict__["register"] = lambda: None
74    return graph_config
75
76
77class FakeParameters(dict):
78    strict = True
79
80
81class FakeOptimization(OptimizationStrategy):
82    def __init__(self, mode, *args, **kwargs):
83        super(FakeOptimization, self).__init__(*args, **kwargs)
84        self.mode = mode
85
86    def should_remove_task(self, task, params, arg):
87        if self.mode == "always":
88            return True
89        if self.mode == "even":
90            return task.task["i"] % 2 == 0
91        if self.mode == "odd":
92            return task.task["i"] % 2 != 0
93        return False
94
95
96class TestGenerator(unittest.TestCase):
97    @pytest.fixture(autouse=True)
98    def patch(self, monkeypatch):
99        self.patch = monkeypatch
100
101    def maketgg(self, target_tasks=None, kinds=[("_fake", [])], params=None):
102        params = params or {}
103        FakeKind.loaded_kinds = []
104        self.target_tasks = target_tasks or []
105
106        def target_tasks_method(full_task_graph, parameters, graph_config):
107            return self.target_tasks
108
109        fake_registry = {
110            mode: FakeOptimization(mode) for mode in ("always", "never", "even", "odd")
111        }
112
113        target_tasks_mod._target_task_methods["test_method"] = target_tasks_method
114        self.patch.setattr(optimize_mod, "registry", fake_registry)
115
116        parameters = FakeParameters(
117            {
118                "_kinds": kinds,
119                "backstop": False,
120                "target_tasks_method": "test_method",
121                "test_manifest_loader": "default",
122                "try_mode": None,
123                "try_task_config": {},
124                "tasks_for": "hg-push",
125                "project": "mozilla-central",
126            }
127        )
128        parameters.update(params)
129
130        self.patch.setattr(generator, "load_graph_config", fake_load_graph_config)
131
132        return WithFakeKind("/root", parameters)
133
134    def test_kind_ordering(self):
135        "When task kinds depend on each other, they are loaded in postorder"
136        self.tgg = self.maketgg(
137            kinds=[
138                ("_fake3", {"kind-dependencies": ["_fake2", "_fake1"]}),
139                ("_fake2", {"kind-dependencies": ["_fake1"]}),
140                ("_fake1", {"kind-dependencies": []}),
141            ]
142        )
143        self.tgg._run_until("full_task_set")
144        self.assertEqual(FakeKind.loaded_kinds, ["_fake1", "_fake2", "_fake3"])
145
146    def test_full_task_set(self):
147        "The full_task_set property has all tasks"
148        self.tgg = self.maketgg()
149        self.assertEqual(
150            self.tgg.full_task_set.graph,
151            graph.Graph({"_fake-t-0", "_fake-t-1", "_fake-t-2"}, set()),
152        )
153        self.assertEqual(
154            sorted(self.tgg.full_task_set.tasks.keys()),
155            sorted(["_fake-t-0", "_fake-t-1", "_fake-t-2"]),
156        )
157
158    def test_full_task_graph(self):
159        "The full_task_graph property has all tasks, and links"
160        self.tgg = self.maketgg()
161        self.assertEqual(
162            self.tgg.full_task_graph.graph,
163            graph.Graph(
164                {"_fake-t-0", "_fake-t-1", "_fake-t-2"},
165                {
166                    ("_fake-t-1", "_fake-t-0", "prev"),
167                    ("_fake-t-2", "_fake-t-1", "prev"),
168                },
169            ),
170        )
171        self.assertEqual(
172            sorted(self.tgg.full_task_graph.tasks.keys()),
173            sorted(["_fake-t-0", "_fake-t-1", "_fake-t-2"]),
174        )
175
176    def test_target_task_set(self):
177        "The target_task_set property has the targeted tasks"
178        self.tgg = self.maketgg(["_fake-t-1"])
179        self.assertEqual(
180            self.tgg.target_task_set.graph, graph.Graph({"_fake-t-1"}, set())
181        )
182        self.assertEqual(
183            set(six.iterkeys(self.tgg.target_task_set.tasks)), {"_fake-t-1"}
184        )
185
186    def test_target_task_graph(self):
187        "The target_task_graph property has the targeted tasks and deps"
188        self.tgg = self.maketgg(["_fake-t-1"])
189        self.assertEqual(
190            self.tgg.target_task_graph.graph,
191            graph.Graph(
192                {"_fake-t-0", "_fake-t-1"}, {("_fake-t-1", "_fake-t-0", "prev")}
193            ),
194        )
195        self.assertEqual(
196            sorted(self.tgg.target_task_graph.tasks.keys()),
197            sorted(["_fake-t-0", "_fake-t-1"]),
198        )
199
200    def test_always_target_tasks(self):
201        "The target_task_graph includes tasks with 'always_target'"
202        tgg_args = {
203            "target_tasks": ["_fake-t-0", "_fake-t-1", "_ignore-t-0", "_ignore-t-1"],
204            "kinds": [
205                ("_fake", {"job-defaults": {"optimization": {"odd": None}}}),
206                (
207                    "_ignore",
208                    {
209                        "job-defaults": {
210                            "attributes": {"always_target": True},
211                            "optimization": {"even": None},
212                        }
213                    },
214                ),
215            ],
216            "params": {"optimize_target_tasks": False},
217        }
218        self.tgg = self.maketgg(**tgg_args)
219        self.assertEqual(
220            sorted(self.tgg.target_task_set.tasks.keys()),
221            sorted(["_fake-t-0", "_fake-t-1", "_ignore-t-0", "_ignore-t-1"]),
222        )
223        self.assertEqual(
224            sorted(self.tgg.target_task_graph.tasks.keys()),
225            sorted(
226                ["_fake-t-0", "_fake-t-1", "_ignore-t-0", "_ignore-t-1", "_ignore-t-2"]
227            ),
228        )
229        self.assertEqual(
230            sorted([t.label for t in self.tgg.optimized_task_graph.tasks.values()]),
231            sorted(["_fake-t-0", "_fake-t-1", "_ignore-t-0", "_ignore-t-1"]),
232        )
233
234    def test_optimized_task_graph(self):
235        "The optimized task graph contains task ids"
236        self.tgg = self.maketgg(["_fake-t-2"])
237        tid = self.tgg.label_to_taskid
238        self.assertEqual(
239            self.tgg.optimized_task_graph.graph,
240            graph.Graph(
241                {tid["_fake-t-0"], tid["_fake-t-1"], tid["_fake-t-2"]},
242                {
243                    (tid["_fake-t-1"], tid["_fake-t-0"], "prev"),
244                    (tid["_fake-t-2"], tid["_fake-t-1"], "prev"),
245                },
246            ),
247        )
248
249
250def test_load_tasks_for_kind(monkeypatch):
251    """
252    `load_tasks_for_kinds` will load the tasks for the provided kind
253    """
254    monkeypatch.setattr(generator, "TaskGraphGenerator", WithFakeKind)
255    monkeypatch.setattr(generator, "load_graph_config", fake_load_graph_config)
256
257    tasks = load_tasks_for_kind(
258        {"_kinds": [("_example-kind", []), ("docker-image", [])]},
259        "_example-kind",
260        "/root",
261    )
262    assert "t-1" in tasks and tasks["t-1"].label == "_example-kind-t-1"
263
264
265if __name__ == "__main__":
266    main()
267