Bug 1659187 - [taskgraph.test] Convert test_optimize.py to the pytest format and use parametrization, r=taskgraph-reviewers,aki
authorAndrew Halberstadt <ahalberstadt@mozilla.com>
Thu, 27 Aug 2020 21:36:14 +0000
changeset 546767 e1c247e7155a715a49da3c9143595b2fa014cc58
parent 546766 f7912b1443a5fb456dc1feadff369da3de5fd5fb
child 546768 e57bccb01008b5f0662df428797e43b074060eb5
push id37737
push usercsabou@mozilla.com
push dateSat, 29 Aug 2020 09:12:26 +0000
treeherdermozilla-central@fdf95334aded [default view] [failures only]
perfherder[talos] [build metrics] [platform microbench] (compared to previous push)
reviewerstaskgraph-reviewers, aki
bugs1659187
milestone82.0a1
first release with
nightly linux32
nightly linux64
nightly mac
nightly win32
nightly win64
last release without
nightly linux32
nightly linux64
nightly mac
nightly win32
nightly win64
Bug 1659187 - [taskgraph.test] Convert test_optimize.py to the pytest format and use parametrization, r=taskgraph-reviewers,aki This removes boiler plate around adding new test cases. Depends on D88482 Differential Revision: https://phabricator.services.mozilla.com/D88483
taskcluster/taskgraph/test/test_optimize.py
--- a/taskcluster/taskgraph/test/test_optimize.py
+++ b/taskcluster/taskgraph/test/test_optimize.py
@@ -1,265 +1,380 @@
 # 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
 
 from functools import partial
-import unittest
 
 import pytest
 from taskgraph import graph, optimize
 from taskgraph.optimize import OptimizationStrategy, All, Any
 from taskgraph.taskgraph import TaskGraph
 from taskgraph.task import Task
 from mozunit import main
-from slugid import nice as slugid
-
-
-@pytest.fixture
-def set_monkeypatch(request, monkeypatch):
-    request.cls.monkeypatch = monkeypatch
 
 
 class Remove(OptimizationStrategy):
 
     def should_remove_task(self, task, params, arg):
         return True
 
 
 class Replace(OptimizationStrategy):
 
     def should_replace_task(self, task, params, taskid):
         return taskid
 
 
-@pytest.mark.usefixtures("set_monkeypatch")
-class TestOptimize(unittest.TestCase):
-
-    strategies = {
+def default_strategies():
+    return {
         'never': OptimizationStrategy(),
         'remove': Remove(),
         'replace': Replace(),
     }
 
-    def make_task(self, label, optimization=None, task_def=None, optimized=None,
-                  task_id=None, dependencies=None):
-        task_def = task_def or {'sample': 'task-def'}
-        task = Task(kind='test', label=label, attributes={}, task=task_def)
-        task.optimization = optimization
-        task.task_id = task_id
-        if dependencies is not None:
-            task.task['dependencies'] = sorted(dependencies)
-        return task
 
-    def make_graph(self, *tasks_and_edges):
-        tasks = {t.label: t for t in tasks_and_edges if isinstance(t, Task)}
-        edges = {e for e in tasks_and_edges if not isinstance(e, Task)}
-        return TaskGraph(tasks, graph.Graph(set(tasks), edges))
+def make_task(
+    label,
+    optimization=None,
+    task_def=None,
+    optimized=None,
+    task_id=None,
+    dependencies=None
+):
+    task_def = task_def or {'sample': 'task-def'}
+    task = Task(kind='test', label=label, attributes={}, task=task_def)
+    task.optimization = optimization
+    task.task_id = task_id
+    if dependencies is not None:
+        task.task['dependencies'] = sorted(dependencies)
+    return task
+
+
+def make_graph(*tasks_and_edges):
+    tasks = {t.label: t for t in tasks_and_edges if isinstance(t, Task)}
+    edges = {e for e in tasks_and_edges if not isinstance(e, Task)}
+    return TaskGraph(tasks, graph.Graph(set(tasks), edges))
+
+
+def make_opt_graph(*tasks_and_edges):
+    tasks = {t.task_id: t for t in tasks_and_edges if isinstance(t, Task)}
+    edges = {e for e in tasks_and_edges if not isinstance(e, Task)}
+    return TaskGraph(tasks, graph.Graph(set(tasks), edges))
 
-    def make_opt_graph(self, *tasks_and_edges):
-        tasks = {t.task_id: t for t in tasks_and_edges if isinstance(t, Task)}
-        edges = {e for e in tasks_and_edges if not isinstance(e, Task)}
-        return TaskGraph(tasks, graph.Graph(set(tasks), edges))
 
-    def make_triangle(self, **opts):
-        """
-        Make a "triangle" graph like this:
+def make_triangle(**opts):
+    """
+    Make a "triangle" graph like this:
 
-          t1 <-------- t3
-           `---- t2 --'
-        """
-        return self.make_graph(
-            self.make_task('t1', opts.get('t1')),
-            self.make_task('t2', opts.get('t2')),
-            self.make_task('t3', opts.get('t3')),
-            ('t3', 't2', 'dep'),
-            ('t3', 't1', 'dep2'),
-            ('t2', 't1', 'dep'))
+      t1 <-------- t3
+       `---- t2 --'
+    """
+    return make_graph(
+        make_task('t1', opts.get('t1')),
+        make_task('t2', opts.get('t2')),
+        make_task('t3', opts.get('t3')),
+        ('t3', 't2', 'dep'),
+        ('t3', 't1', 'dep2'),
+        ('t2', 't1', 'dep'),
+    )
+
+
+@pytest.mark.parametrize(
+    "graph,kwargs,exp_removed",
+    (
+        # A graph full of optimization=never has nothing removed
+        pytest.param(
+            make_triangle(),
+            {},
+            # expectations
+            set(),
+            id="never",
+        ),
 
-    def assert_remove_tasks(self, graph, exp_removed, do_not_optimize=set(), strategies=None):
-        strategies = strategies or self.strategies
-        got_removed = optimize.remove_tasks(
-            target_task_graph=graph,
-            optimizations=optimize._get_optimizations(graph, strategies),
-            params={},
-            do_not_optimize=do_not_optimize,
-            requested_tasks=graph)
-        self.assertEqual(got_removed, exp_removed)
+        # A graph full of optimization=remove removes everything
+        pytest.param(
+            make_triangle(
+                t1={'remove': None},
+                t2={'remove': None},
+                t3={'remove': None},
+            ),
+            {},
+            # expectations
+            {"t1", "t2", "t3"},
+            id="all",
+        ),
 
-    def test_remove_tasks_never(self):
-        "A graph full of optimization=never has nothing removed"
-        graph = self.make_triangle()
-        self.assert_remove_tasks(graph, set())
-
-    def test_remove_tasks_all(self):
-        "A graph full of optimization=remove has removes everything"
-        graph = self.make_triangle(
-            t1={'remove': None},
-            t2={'remove': None},
-            t3={'remove': None})
-        self.assert_remove_tasks(graph, {'t1', 't2', 't3'})
+        # Tasks with the 'any' composite strategy are removed when any substrategy says to
+        pytest.param(
+            make_triangle(
+                t1={'any': None},
+                t2={'any': None},
+                t3={'any': None},
+            ),
+            {"strategies": lambda: {"any": Any("never", "remove")}},
+            # expectations
+            {"t1", "t2", "t3"},
+            id="composite_strategies_any",
+        ),
 
-    def test_composite_strategies_any(self):
-        self.monkeypatch.setattr(optimize, 'registry', self.strategies)
-        strategies = self.strategies.copy()
-        strategies['any'] = Any('never', 'remove')
-
-        graph = self.make_triangle(
-            t1={'any': None},
-            t2={'any': None},
-            t3={'any': None})
-
-        self.assert_remove_tasks(graph, {'t1', 't2', 't3'}, strategies=strategies)
+        # Tasks with the 'all' composite strategy are removed when all substrategies say to
+        pytest.param(
+            make_triangle(
+                t1={'all': None},
+                t2={'all': None},
+                t3={'all': None},
+            ),
+            {"strategies": lambda: {"all": All("never", "remove")}},
+            # expectations
+            set(),
+            id="composite_strategies_all",
+        ),
 
-    def test_composite_strategies_all(self):
-        self.monkeypatch.setattr(optimize, 'registry', self.strategies)
-        strategies = self.strategies.copy()
-        strategies['all'] = All('never', 'remove')
+        # Removable tasks that are depended on by non-removable tasks are not removed
+        pytest.param(
+            make_triangle(
+                t1={'remove': None},
+                t3={'remove': None},
+            ),
+            {},
+            # expectations
+            {"t3"},
+            id="blocked",
+        ),
 
-        graph = self.make_triangle(
-            t1={'all': None},
-            t2={'all': None},
-            t3={'all': None})
-        self.assert_remove_tasks(graph, set(), strategies=strategies)
+        # Removable tasks that are marked do_not_optimize are not removed
+        pytest.param(
+            make_triangle(
+                t1={'remove': None},
+                t2={'remove': None},  # but do_not_optimize
+                t3={'remove': None},
+            ),
+            {"do_not_optimize": {"t2"}},
+            # expectations
+            {"t3"},
+            id="do_not_optimize",
+        ),
+    )
+)
+def test_remove_tasks(monkeypatch, graph, kwargs, exp_removed):
+    """Tests the `remove_tasks` function.
 
-    def test_remove_tasks_blocked(self):
-        "Removable tasks that are depended on by non-removable tasks are not removed"
-        graph = self.make_triangle(
-            t1={'remove': None},
-            t3={'remove': None})
-        self.assert_remove_tasks(graph, {'t3'})
+    Each test case takes three arguments:
 
-    def test_remove_tasks_do_not_optimize(self):
-        "Removable tasks that are marked do_not_optimize are not removed"
-        graph = self.make_triangle(
-            t1={'remove': None},
-            t2={'remove': None},  # but do_not_optimize
-            t3={'remove': None})
-        self.assert_remove_tasks(graph, {'t3'}, do_not_optimize={'t2'})
+    1. A `TaskGraph` instance.
+    2. Keyword arguments to pass into `remove_tasks`.
+    3. The set of task labels that are expected to be removed.
+    """
+    # set up strategies
+    strategies = default_strategies()
+    monkeypatch.setattr(optimize, "registry", strategies)
+    strategies = kwargs.pop('strategies', None) or strategies
+    if callable(strategies):
+        strategies = strategies()
 
-    def assert_replace_tasks(self, graph, exp_replaced, exp_removed=set(), exp_label_to_taskid={},
-                             do_not_optimize=None, label_to_taskid=None, removed_tasks=None,
-                             existing_tasks=None):
-        do_not_optimize = do_not_optimize or set()
-        label_to_taskid = label_to_taskid or {}
-        removed_tasks = removed_tasks or set()
-        existing_tasks = existing_tasks or {}
+    kwargs.setdefault("params", {})
+    kwargs.setdefault("do_not_optimize", set())
+    kwargs.setdefault("requested_tasks", graph)
+
+    got_removed = optimize.remove_tasks(
+        target_task_graph=graph,
+        optimizations=optimize._get_optimizations(graph, strategies),
+        **kwargs
+    )
+    assert got_removed == exp_removed
+
 
-        got_replaced = optimize.replace_tasks(
-            target_task_graph=graph,
-            optimizations=optimize._get_optimizations(graph, self.strategies),
-            params={},
-            do_not_optimize=do_not_optimize,
-            label_to_taskid=label_to_taskid,
-            removed_tasks=removed_tasks,
-            existing_tasks=existing_tasks)
-        self.assertEqual(got_replaced, exp_replaced)
-        self.assertEqual(removed_tasks, exp_removed)
-        self.assertEqual(label_to_taskid, exp_label_to_taskid)
+@pytest.mark.parametrize(
+    "graph,kwargs,exp_replaced,exp_removed,exp_label_to_taskid",
+    (
+        # A task cannot be replaced if it depends on one that was not replaced
+        pytest.param(
+            make_triangle(
+                t1={'replace': 'e1'},
+                t3={'replace': 'e3'},
+            ),
+            {},
+            # expectations
+            {'t1'},
+            set(),
+            {'t1': 'e1'},
+            id="blocked",
+        ),
 
-    def test_replace_tasks_never(self):
-        "No tasks are replaced when strategy is 'never'"
-        graph = self.make_triangle()
-        self.assert_replace_tasks(graph, set())
+        # A task cannot be replaced if it should not be optimized
+        pytest.param(
+            make_triangle(
+                t1={'replace': 'e1'},
+                t2={'replace': 'xxx'},  # but do_not_optimize
+                t3={'replace': 'e3'},
+            ),
+            {"do_not_optimize": {"t2"}},
+            # expectations
+            {'t1'},
+            set(),
+            {'t1': 'e1'},
+            id="do_not_optimize",
+        ),
 
-    def test_replace_tasks_all(self):
-        "All replacable tasks are replaced when strategy is 'replace'"
-        graph = self.make_triangle(
-            t1={'replace': 'e1'},
-            t2={'replace': 'e2'},
-            t3={'replace': 'e3'})
-        self.assert_replace_tasks(
-            graph,
-            exp_replaced={'t1', 't2', 't3'},
-            exp_label_to_taskid={'t1': 'e1', 't2': 'e2', 't3': 'e3'})
+        # No tasks are replaced when strategy is 'never'
+        pytest.param(
+            make_triangle(),
+            {},
+            set(),
+            set(),
+            {},
+            id="never",
+        ),
 
-    def test_replace_tasks_blocked(self):
-        "A task cannot be replaced if it depends on one that was not replaced"
-        graph = self.make_triangle(
-            t1={'replace': 'e1'},
-            t3={'replace': 'e3'})
-        self.assert_replace_tasks(
-            graph,
-            exp_replaced={'t1'},
-            exp_label_to_taskid={'t1': 'e1'})
-
-    def test_replace_tasks_do_not_optimize(self):
-        "A task cannot be replaced if it depends on one that was not replaced"
-        graph = self.make_triangle(
-            t1={'replace': 'e1'},
-            t2={'replace': 'xxx'},  # but do_not_optimize
-            t3={'replace': 'e3'})
-        self.assert_replace_tasks(
-            graph,
-            exp_replaced={'t1'},
-            exp_label_to_taskid={'t1': 'e1'},
-            do_not_optimize={'t2'})
+        # All replacable tasks are replaced when strategy is 'replace'
+        pytest.param(
+            make_triangle(
+                t1={'replace': 'e1'},
+                t2={'replace': 'e2'},
+                t3={'replace': 'e3'},
+            ),
+            {},
+            # expectations
+            {'t1', 't2', 't3'},
+            set(),
+            {'t1': 'e1', 't2': 'e2', 't3': 'e3'},
+            id="all",
+        ),
 
-    def test_replace_tasks_removed(self):
-        "A task can be replaced with nothing"
-        graph = self.make_triangle(
-            t1={'replace': 'e1'},
-            t2={'replace': True},
-            t3={'replace': True})
-        self.assert_replace_tasks(
-            graph,
-            exp_replaced={'t1'},
-            exp_removed={'t2', 't3'},
-            exp_label_to_taskid={'t1': 'e1'})
+        # A task can be replaced with nothing
+        pytest.param(
+            make_triangle(
+                t1={'replace': 'e1'},
+                t2={'replace': True},
+                t3={'replace': True},
+            ),
+            {},
+            # expectations
+            {'t1'},
+            {'t2', 't3'},
+            {'t1': 'e1'},
+            id="tasks_removed",
+        ),
+    )
+)
+def test_replace_tasks(
+    graph,
+    kwargs,
+    exp_replaced,
+    exp_removed,
+    exp_label_to_taskid,
+):
+    """Tests the `replace_tasks` function.
+
+    Each test case takes five arguments:
 
-    def assert_subgraph(self, graph, removed_tasks, replaced_tasks,
-                        label_to_taskid, exp_subgraph, exp_label_to_taskid):
-        self.maxDiff = None
-        optimize.slugid = partial(next, (b'tid%d' % i for i in range(1, 10)))
-        try:
-            got_subgraph = optimize.get_subgraph(graph, removed_tasks,
-                                                 replaced_tasks, label_to_taskid,
-                                                 "DECISION-TASK")
-        finally:
-            optimize.slugid = slugid
-        self.assertEqual(got_subgraph.graph, exp_subgraph.graph)
-        self.assertEqual(got_subgraph.tasks, exp_subgraph.tasks)
-        self.assertEqual(label_to_taskid, exp_label_to_taskid)
+    1. A `TaskGraph` instance.
+    2. Keyword arguments to pass into `replace_tasks`.
+    3. The set of task labels that are expected to be replaced.
+    4. The set of task labels that are expected to be removed.
+    5. The expected label_to_taskid.
+    """
+    kwargs.setdefault("params", {})
+    kwargs.setdefault("do_not_optimize", set())
+    kwargs.setdefault("label_to_taskid", {})
+    kwargs.setdefault("removed_tasks", set())
+    kwargs.setdefault("existing_tasks", {})
 
-    def test_get_subgraph_no_change(self):
-        "get_subgraph returns a similarly-shaped subgraph when nothing is removed"
-        graph = self.make_triangle()
-        self.assert_subgraph(
-            graph, set(), set(), {},
-            self.make_opt_graph(
-                self.make_task('t1', task_id='tid1', dependencies={}),
-                self.make_task('t2', task_id='tid2', dependencies={'tid1'}),
-                self.make_task('t3', task_id='tid3', dependencies={'tid1', 'tid2'}),
+    got_replaced = optimize.replace_tasks(
+        target_task_graph=graph,
+        optimizations=optimize._get_optimizations(graph, default_strategies()),
+        **kwargs
+    )
+    assert got_replaced == exp_replaced
+    assert kwargs["removed_tasks"] == exp_removed
+    assert kwargs["label_to_taskid"] == exp_label_to_taskid
+
+
+@pytest.mark.parametrize(
+    "graph,kwargs,exp_subgraph,exp_label_to_taskid",
+    (
+        # Test get_subgraph returns a similarly-shaped subgraph when nothing is removed
+        pytest.param(
+            make_triangle(),
+            {},
+            make_opt_graph(
+                make_task('t1', task_id='tid1', dependencies={}),
+                make_task('t2', task_id='tid2', dependencies={'tid1'}),
+                make_task('t3', task_id='tid3', dependencies={'tid1', 'tid2'}),
                 ('tid3', 'tid2', 'dep'),
                 ('tid3', 'tid1', 'dep2'),
-                ('tid2', 'tid1', 'dep')),
-            {'t1': 'tid1', 't2': 'tid2', 't3': 'tid3'})
+                ('tid2', 'tid1', 'dep')
+            ),
+            {'t1': 'tid1', 't2': 'tid2', 't3': 'tid3'},
+            id="no_change",
+        ),
 
-    def test_get_subgraph_removed(self):
-        "get_subgraph returns a smaller subgraph when tasks are removed"
-        graph = self.make_triangle()
-        self.assert_subgraph(
-            graph, {'t2', 't3'}, set(), {},
-            self.make_opt_graph(
-                self.make_task('t1', task_id='tid1', dependencies={})),
-            {'t1': 'tid1'})
+        # Test get_subgraph returns a smaller subgraph when tasks are removed
+        pytest.param(
+            make_triangle(),
+            {
+                "removed_tasks": {"t2", "t3"},
+            },
+            # expectations
+            make_opt_graph(
+                make_task('t1', task_id='tid1', dependencies={})
+            ),
+            {'t1': 'tid1'},
+            id="removed",
+        ),
 
-    def test_get_subgraph_replaced(self):
-        "get_subgraph returns a smaller subgraph when tasks are replaced"
-        graph = self.make_triangle()
-        self.assert_subgraph(
-            graph, set(), {'t1', 't2'}, {'t1': 'e1', 't2': 'e2'},
-            self.make_opt_graph(
-                self.make_task('t3', task_id='tid1', dependencies={'e1', 'e2'})),
-            {'t1': 'e1', 't2': 'e2', 't3': 'tid1'})
+        # Test get_subgraph returns a smaller subgraph when tasks are replaced
+        pytest.param(
+            make_triangle(),
+            {
+                "replaced_tasks": {"t1", "t2"},
+                "label_to_taskid": {"t1": "e1", "t2": "e2"},
+            },
+            # expectations
+            make_opt_graph(
+                make_task('t3', task_id='tid1', dependencies={'e1', 'e2'})
+            ),
+            {'t1': 'e1', 't2': 'e2', 't3': 'tid1'},
+            id="replaced",
+        ),
+    )
+)
+def test_get_subgraph(
+    monkeypatch,
+    graph,
+    kwargs,
+    exp_subgraph,
+    exp_label_to_taskid
+):
+    """Tests the `get_subgraph` function.
 
-    def test_get_subgraph_removed_dep(self):
-        "get_subgraph raises an Exception when a task depends on a removed task"
-        graph = self.make_triangle()
-        with self.assertRaises(Exception):
-            optimize.get_subgraph(graph, {'t2'}, set(), {})
+    Each test case takes 4 arguments:
+
+    1. A `TaskGraph` instance.
+    2. Keyword arguments to pass into `get_subgraph`.
+    3. The expected subgraph.
+    4. The expected label_to_taskid.
+    """
+    monkeypatch.setattr(optimize, "slugid", partial(next, (b'tid%d' % i for i in range(1, 10))))
+
+    kwargs.setdefault("removed_tasks", set())
+    kwargs.setdefault("replaced_tasks", set())
+    kwargs.setdefault("label_to_taskid", {})
+    kwargs.setdefault("decision_task_id", "DECISION-TASK")
+
+    got_subgraph = optimize.get_subgraph(graph, **kwargs)
+    assert got_subgraph.graph == exp_subgraph.graph
+    assert got_subgraph.tasks == exp_subgraph.tasks
+    assert kwargs["label_to_taskid"] == exp_label_to_taskid
+
+
+def test_get_subgraph_removed_dep():
+    "get_subgraph raises an Exception when a task depends on a removed task"
+    graph = make_triangle()
+    with pytest.raises(Exception):
+        optimize.get_subgraph(graph, {'t2'}, set(), {})
 
 
 if __name__ == '__main__':
     main()