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 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(),
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)