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 test_remove_tasks(monkeypatch, graph, kwargs, exp_removed):
"""Tests the `remove_tasks` function.
Each test case takes three arguments:
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)
extra = kwargs.pop("strategies", None)
if extra:
if callable(extra):
extra = extra()
strategies.update(extra)
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
def assert_remove_tasks(self, graph, exp_removed, do_not_optimize=set()):
got_removed = optimize.remove_tasks(
target_task_graph=graph,
optimizations=optimize._get_optimizations(graph, self.strategies),
params={},
do_not_optimize=do_not_optimize)
self.assertEqual(got_removed, exp_removed)
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_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 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)
