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
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 {}
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)
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 {}
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)