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 unittest
9from mozunit import main
10
11from taskgraph.generator import TaskGraphGenerator, Kind
12from taskgraph.optimize import OptimizationStrategy
13from taskgraph.util.templates import merge
14from taskgraph import (
15    generator,
16    graph,
17    optimize as optimize_mod,
18    target_tasks as target_tasks_mod,
19)
20
21
22def fake_loader(kind, path, config, parameters, loaded_tasks):
23    for i in range(3):
24        dependencies = {}
25        if i >= 1:
26            dependencies['prev'] = '{}-t-{}'.format(kind, i-1)
27
28        task = {
29            'kind': kind,
30            'label': '{}-t-{}'.format(kind, i),
31            'attributes': {'_tasknum': str(i)},
32            'task': {'i': i},
33            'dependencies': dependencies,
34        }
35        if 'job-defaults' in config:
36            task = merge(config['job-defaults'], task)
37        yield task
38
39
40class FakeKind(Kind):
41
42    def _get_loader(self):
43        return fake_loader
44
45    def load_tasks(self, parameters, loaded_tasks):
46        FakeKind.loaded_kinds.append(self.name)
47        return super(FakeKind, self).load_tasks(parameters, loaded_tasks)
48
49
50class WithFakeKind(TaskGraphGenerator):
51
52    def _load_kinds(self, graph_config):
53        for kind_name, cfg in self.parameters['_kinds']:
54            config = {
55                'transforms': [],
56            }
57            if cfg:
58                config.update(cfg)
59            yield FakeKind(kind_name, '/fake', config, graph_config)
60
61
62def fake_load_graph_config(root_dir):
63    return {'trust-domain': 'test-domain'}
64
65
66class FakeParameters(dict):
67    strict = True
68
69
70class FakeOptimization(OptimizationStrategy):
71    def __init__(self, mode, *args, **kwargs):
72        super(FakeOptimization, self).__init__(*args, **kwargs)
73        self.mode = mode
74
75    def should_remove_task(self, task, params, arg):
76        if self.mode == 'always':
77            return True
78        if self.mode == 'even':
79            return task.task['i'] % 2 == 0
80        if self.mode == 'odd':
81            return task.task['i'] % 2 != 0
82        return False
83
84
85class TestGenerator(unittest.TestCase):
86
87    @pytest.fixture(autouse=True)
88    def patch(self, monkeypatch):
89        self.patch = monkeypatch
90
91    def maketgg(self, target_tasks=None, kinds=[('_fake', [])], params=None):
92        params = params or {}
93        FakeKind.loaded_kinds = []
94        self.target_tasks = target_tasks or []
95
96        def target_tasks_method(full_task_graph, parameters, graph_config):
97            return self.target_tasks
98
99        def make_fake_strategies():
100            return {mode: FakeOptimization(mode)
101                    for mode in ('always', 'never', 'even', 'odd')}
102
103        target_tasks_mod._target_task_methods['test_method'] = target_tasks_method
104        self.patch.setattr(optimize_mod, '_make_default_strategies', make_fake_strategies)
105
106        parameters = FakeParameters({
107            '_kinds': kinds,
108            'target_tasks_method': 'test_method',
109            'try_mode': None,
110        })
111        parameters.update(params)
112
113        self.patch.setattr(generator, 'load_graph_config', fake_load_graph_config)
114
115        return WithFakeKind('/root', parameters)
116
117    def test_kind_ordering(self):
118        "When task kinds depend on each other, they are loaded in postorder"
119        self.tgg = self.maketgg(kinds=[
120            ('_fake3', {'kind-dependencies': ['_fake2', '_fake1']}),
121            ('_fake2', {'kind-dependencies': ['_fake1']}),
122            ('_fake1', {'kind-dependencies': []}),
123        ])
124        self.tgg._run_until('full_task_set')
125        self.assertEqual(FakeKind.loaded_kinds, ['_fake1', '_fake2', '_fake3'])
126
127    def test_full_task_set(self):
128        "The full_task_set property has all tasks"
129        self.tgg = self.maketgg()
130        self.assertEqual(self.tgg.full_task_set.graph,
131                         graph.Graph({'_fake-t-0', '_fake-t-1', '_fake-t-2'}, set()))
132        self.assertEqual(sorted(self.tgg.full_task_set.tasks.keys()),
133                         sorted(['_fake-t-0', '_fake-t-1', '_fake-t-2']))
134
135    def test_full_task_graph(self):
136        "The full_task_graph property has all tasks, and links"
137        self.tgg = self.maketgg()
138        self.assertEqual(self.tgg.full_task_graph.graph,
139                         graph.Graph({'_fake-t-0', '_fake-t-1', '_fake-t-2'},
140                                     {
141                                         ('_fake-t-1', '_fake-t-0', 'prev'),
142                                         ('_fake-t-2', '_fake-t-1', 'prev'),
143                         }))
144        self.assertEqual(sorted(self.tgg.full_task_graph.tasks.keys()),
145                         sorted(['_fake-t-0', '_fake-t-1', '_fake-t-2']))
146
147    def test_target_task_set(self):
148        "The target_task_set property has the targeted tasks"
149        self.tgg = self.maketgg(['_fake-t-1'])
150        self.assertEqual(self.tgg.target_task_set.graph,
151                         graph.Graph({'_fake-t-1'}, set()))
152        self.assertEqual(self.tgg.target_task_set.tasks.keys(),
153                         ['_fake-t-1'])
154
155    def test_target_task_graph(self):
156        "The target_task_graph property has the targeted tasks and deps"
157        self.tgg = self.maketgg(['_fake-t-1'])
158        self.assertEqual(self.tgg.target_task_graph.graph,
159                         graph.Graph({'_fake-t-0', '_fake-t-1'},
160                                     {('_fake-t-1', '_fake-t-0', 'prev')}))
161        self.assertEqual(sorted(self.tgg.target_task_graph.tasks.keys()),
162                         sorted(['_fake-t-0', '_fake-t-1']))
163
164    def test_always_target_tasks(self):
165        "The target_task_graph includes tasks with 'always_target'"
166        tgg_args = {
167            'target_tasks': ['_fake-t-0', '_fake-t-1', '_ignore-t-0', '_ignore-t-1'],
168            'kinds': [
169                ('_fake', {'job-defaults': {'optimization': {'odd': None}}}),
170                ('_ignore', {'job-defaults': {
171                    'attributes': {'always_target': True},
172                    'optimization': {'even': None},
173                }}),
174            ],
175            'params': {'optimize_target_tasks': False},
176        }
177        self.tgg = self.maketgg(**tgg_args)
178        self.assertEqual(
179            sorted(self.tgg.target_task_set.tasks.keys()),
180            sorted(['_fake-t-0', '_fake-t-1', '_ignore-t-0', '_ignore-t-1']))
181        self.assertEqual(
182            sorted(self.tgg.target_task_graph.tasks.keys()),
183            sorted(['_fake-t-0', '_fake-t-1', '_ignore-t-0', '_ignore-t-1', '_ignore-t-2']))
184        self.assertEqual(
185            sorted([t.label for t in self.tgg.optimized_task_graph.tasks.values()]),
186            sorted(['_fake-t-0', '_fake-t-1', '_ignore-t-0', '_ignore-t-1']))
187
188    def test_optimized_task_graph(self):
189        "The optimized task graph contains task ids"
190        self.tgg = self.maketgg(['_fake-t-2'])
191        tid = self.tgg.label_to_taskid
192        self.assertEqual(
193            self.tgg.optimized_task_graph.graph,
194            graph.Graph({tid['_fake-t-0'], tid['_fake-t-1'], tid['_fake-t-2']}, {
195                (tid['_fake-t-1'], tid['_fake-t-0'], 'prev'),
196                (tid['_fake-t-2'], tid['_fake-t-1'], 'prev'),
197            }))
198
199
200if __name__ == '__main__':
201    main()
202