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
This commit is contained in:
Andrew Halberstadt 2020-08-27 21:36:14 +00:00
Родитель b84d7787f5
Коммит 59d997abe4
1 изменённых файлов: 315 добавлений и 200 удалений

Просмотреть файл

@ -5,7 +5,6 @@
from __future__ import absolute_import, print_function, unicode_literals
from functools import partial
import unittest
import pytest
from taskgraph import graph, optimize
@ -13,12 +12,6 @@ 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):
@ -33,232 +26,354 @@ class Replace(OptimizationStrategy):
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_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_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))
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'))
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)
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 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'})
def make_triangle(**opts):
"""
Make a "triangle" graph like this:
def test_composite_strategies_any(self):
self.monkeypatch.setattr(optimize, 'registry', self.strategies)
strategies = self.strategies.copy()
strategies['any'] = Any('never', 'remove')
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'),
)
graph = self.make_triangle(
t1={'any': None},
t2={'any': None},
t3={'any': None})
self.assert_remove_tasks(graph, {'t1', 't2', 't3'}, strategies=strategies)
@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 test_composite_strategies_all(self):
self.monkeypatch.setattr(optimize, 'registry', self.strategies)
strategies = self.strategies.copy()
strategies['all'] = All('never', 'remove')
# 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",
),
graph = self.make_triangle(
t1={'all': None},
t2={'all': None},
t3={'all': None})
self.assert_remove_tasks(graph, set(), strategies=strategies)
# 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_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'})
# 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_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'})
# 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",
),
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 {}
# 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.
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)
Each test case takes three arguments:
def test_replace_tasks_never(self):
"No tasks are replaced when strategy is 'never'"
graph = self.make_triangle()
self.assert_replace_tasks(graph, set())
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 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'})
kwargs.setdefault("params", {})
kwargs.setdefault("do_not_optimize", set())
kwargs.setdefault("requested_tasks", graph)
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'})
got_removed = optimize.remove_tasks(
target_task_graph=graph,
optimizations=optimize._get_optimizations(graph, strategies),
**kwargs
)
assert got_removed == exp_removed
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'})
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'})
@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 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)
# 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_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'}),
# No tasks are replaced when strategy is 'never'
pytest.param(
make_triangle(),
{},
set(),
set(),
{},
id="never",
),
# 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",
),
# 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:
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", {})
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__':