test_generator.py

# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.

from __future__ import absolute_import, print_function, unicode_literals

import pytest
import unittest
from mozunit import main

from taskgraph.generator import TaskGraphGenerator, Kind
from taskgraph.optimize import OptimizationStrategy
from taskgraph.util.templates import merge
from taskgraph import (
    generator,
    graph,
    optimize as optimize_mod,
    target_tasks as target_tasks_mod,
)


def fake_loader(kind, path, config, parameters, loaded_tasks):
    for i in range(3):
        dependencies = {}
        if i >= 1:
            dependencies['prev'] = '{}-t-{}'.format(kind, i-1)

        task = {
            'kind': kind,
            'label': '{}-t-{}'.format(kind, i),
            'attributes': {'_tasknum': str(i)},
            'task': {'i': i},
            'dependencies': dependencies,
        }
        if 'job-defaults' in config:
            task = merge(config['job-defaults'], task)
        yield task


class FakeKind(Kind):

    def _get_loader(self):
        return fake_loader

    def load_tasks(self, parameters, loaded_tasks):
        FakeKind.loaded_kinds.append(self.name)
        return super(FakeKind, self).load_tasks(parameters, loaded_tasks)


class WithFakeKind(TaskGraphGenerator):

    def _load_kinds(self, graph_config):
        for kind_name, cfg in self.parameters['_kinds']:
            config = {
                'transforms': [],
            }
            if cfg:
                config.update(cfg)
            yield FakeKind(kind_name, '/fake', config, graph_config)


def fake_load_graph_config(root_dir):
    return {'trust-domain': 'test-domain'}


class FakeParameters(dict):
    strict = True


class FakeOptimization(OptimizationStrategy):
    def __init__(self, mode, *args, **kwargs):
        super(FakeOptimization, self).__init__(*args, **kwargs)
        self.mode = mode

    def should_remove_task(self, task, params, arg):
        if self.mode == 'always':
            return True
        if self.mode == 'even':
            return task.task['i'] % 2 == 0
        if self.mode == 'odd':
            return task.task['i'] % 2 != 0
        return False


class TestGenerator(unittest.TestCase):

    @pytest.fixture(autouse=True)
    def patch(self, monkeypatch):
        self.patch = monkeypatch

    def maketgg(self, target_tasks=None, kinds=[('_fake', [])], params=None):
        params = params or {}
        FakeKind.loaded_kinds = []
        self.target_tasks = target_tasks or []

        def target_tasks_method(full_task_graph, parameters, graph_config):
            return self.target_tasks

        def make_fake_strategies():
            return {mode: FakeOptimization(mode)
                    for mode in ('always', 'never', 'even', 'odd')}

        target_tasks_mod._target_task_methods['test_method'] = target_tasks_method
        self.patch.setattr(optimize_mod, '_make_default_strategies', make_fake_strategies)

        parameters = FakeParameters({
            '_kinds': kinds,
            'target_tasks_method': 'test_method',
            'try_mode': None,
        })
        parameters.update(params)

        self.patch.setattr(generator, 'load_graph_config', fake_load_graph_config)

        return WithFakeKind('/root', parameters)

    def test_kind_ordering(self):
        "When task kinds depend on each other, they are loaded in postorder"
        self.tgg = self.maketgg(kinds=[
            ('_fake3', {'kind-dependencies': ['_fake2', '_fake1']}),
            ('_fake2', {'kind-dependencies': ['_fake1']}),
            ('_fake1', {'kind-dependencies': []}),
        ])
        self.tgg._run_until('full_task_set')
        self.assertEqual(FakeKind.loaded_kinds, ['_fake1', '_fake2', '_fake3'])

    def test_full_task_set(self):
        "The full_task_set property has all tasks"
        self.tgg = self.maketgg()
        self.assertEqual(self.tgg.full_task_set.graph,
                         graph.Graph({'_fake-t-0', '_fake-t-1', '_fake-t-2'}, set()))
        self.assertEqual(sorted(self.tgg.full_task_set.tasks.keys()),
                         sorted(['_fake-t-0', '_fake-t-1', '_fake-t-2']))

    def test_full_task_graph(self):
        "The full_task_graph property has all tasks, and links"
        self.tgg = self.maketgg()
        self.assertEqual(self.tgg.full_task_graph.graph,
                         graph.Graph({'_fake-t-0', '_fake-t-1', '_fake-t-2'},
                                     {
                                         ('_fake-t-1', '_fake-t-0', 'prev'),
                                         ('_fake-t-2', '_fake-t-1', 'prev'),
                         }))
        self.assertEqual(sorted(self.tgg.full_task_graph.tasks.keys()),
                         sorted(['_fake-t-0', '_fake-t-1', '_fake-t-2']))

    def test_target_task_set(self):
        "The target_task_set property has the targeted tasks"
        self.tgg = self.maketgg(['_fake-t-1'])
        self.assertEqual(self.tgg.target_task_set.graph,
                         graph.Graph({'_fake-t-1'}, set()))
        self.assertEqual(self.tgg.target_task_set.tasks.keys(),
                         ['_fake-t-1'])

    def test_target_task_graph(self):
        "The target_task_graph property has the targeted tasks and deps"
        self.tgg = self.maketgg(['_fake-t-1'])
        self.assertEqual(self.tgg.target_task_graph.graph,
                         graph.Graph({'_fake-t-0', '_fake-t-1'},
                                     {('_fake-t-1', '_fake-t-0', 'prev')}))
        self.assertEqual(sorted(self.tgg.target_task_graph.tasks.keys()),
                         sorted(['_fake-t-0', '_fake-t-1']))

    def test_always_target_tasks(self):
        "The target_task_graph includes tasks with 'always_target'"
        tgg_args = {
            'target_tasks': ['_fake-t-0', '_fake-t-1', '_ignore-t-0', '_ignore-t-1'],
            'kinds': [
                ('_fake', {'job-defaults': {'optimization': {'odd': None}}}),
                ('_ignore', {'job-defaults': {
                    'attributes': {'always_target': True},
                    'optimization': {'even': None},
                }}),
            ],
            'params': {'optimize_target_tasks': False},
        }
        self.tgg = self.maketgg(**tgg_args)
        self.assertEqual(
            sorted(self.tgg.target_task_set.tasks.keys()),
            sorted(['_fake-t-0', '_fake-t-1', '_ignore-t-0', '_ignore-t-1']))
        self.assertEqual(
            sorted(self.tgg.target_task_graph.tasks.keys()),
            sorted(['_fake-t-0', '_fake-t-1', '_ignore-t-0', '_ignore-t-1', '_ignore-t-2']))
        self.assertEqual(
            sorted([t.label for t in self.tgg.optimized_task_graph.tasks.values()]),
            sorted(['_fake-t-0', '_fake-t-1', '_ignore-t-0', '_ignore-t-1']))

    def test_optimized_task_graph(self):
        "The optimized task graph contains task ids"
        self.tgg = self.maketgg(['_fake-t-2'])
        tid = self.tgg.label_to_taskid
        self.assertEqual(
            self.tgg.optimized_task_graph.graph,
            graph.Graph({tid['_fake-t-0'], tid['_fake-t-1'], tid['_fake-t-2']}, {
                (tid['_fake-t-1'], tid['_fake-t-0'], 'prev'),
                (tid['_fake-t-2'], tid['_fake-t-1'], 'prev'),
            }))


if __name__ == '__main__':
    main()