def test_reset_reference_tasks_to_terminal_tasks():
with Flow(name="test") as f:
t1 = Task()
t2 = Task()
t3 = Task()
f.add_edge(t1, t2)
f.add_edge(t2, t3)
f.set_reference_tasks([t2])
assert f.reference_tasks() == set([t2])
f.set_reference_tasks([])
assert f.reference_tasks() == f.terminal_tasks()
def test_sorted_tasks_with_invalid_start_task():
"""
t1 -> t2 -> t3 -> t4
t3 -> t5
"""
f = Flow(name="test")
t1 = Task("1")
t2 = Task("2")
t3 = Task("3")
f.add_edge(t1, t2)
with pytest.raises(ValueError) as exc:
f.sorted_tasks(root_tasks=[t3])
assert "not found in Flow" in str(exc.value)
def test_cache_all_upstream_edges(self):
f = Flow(name="test")
t1 = Task()
t2 = Task()
t3 = Task()
f.add_edge(t1, t2)
f.all_upstream_edges()
key = ("all_upstream_edges", ())
f._cache[key] = 1
assert f.all_upstream_edges() == 1
f.add_edge(t2, t3)
assert f.all_upstream_edges() != 1
def test_merge():
f1 = Flow(name="test")
f2 = Flow(name="test")
t1 = Task()
t2 = Task()
t3 = Task()
f1.add_edge(t1, t2)
f2.add_edge(t2, t3)
f2.update(f1)
assert f2.tasks == set([t1, t2, t3])
assert len(f2.edges) == 2
def test_cache_survives_pickling(self):
f = Flow(name="test")
t1 = Task()
t2 = Task()
t3 = Task()
f.add_edge(t1, t2)
f.sorted_tasks()
key = ("_sorted_tasks", (("root_tasks", ()), ))
f._cache[key] = 1
assert f.sorted_tasks() == 1
f2 = cloudpickle.loads(cloudpickle.dumps(f))
assert f2.sorted_tasks() == 1
f2.add_edge(t2, t3)
assert f2.sorted_tasks() != 1
def test_equality_based_on_edges(self):
f1 = Flow(name="test")
f2 = Flow(name="test")
t1 = Task()
t2 = Task()
t3 = Task()
for f in [f1, f2]:
f.add_edge(t1, t2)
f.add_edge(t1, t3)
assert f1 == f2
f2.add_edge(t2, t3)
assert f1 != f2
def test_get_tasks_can_check_types(self):
class Specific(Task):
pass
t1, t2 = Task(name="t1", tags=["a", "b"]), Specific(name="t1", tags=["a"])
f = Flow(name="test", tasks=[t1, t2])
assert f.get_tasks(task_type=Specific) == [t2]
def test_viz_renders_if_ipython_isnt_installed_or_errors(self, error):
graphviz = MagicMock()
ipython = MagicMock(get_ipython=MagicMock(side_effect=error))
with patch.dict("sys.modules", graphviz=graphviz, IPython=ipython):
with Flow(name="test") as f:
res = AddTask(name="a_nice_task").map(x=Task(name="a_list_task"), y=8)
f.visualize()
def test_create_flow_with_edges(self):
f1 = Flow(
name="test",
edges=[Edge(upstream_task=Task(), downstream_task=AddTask(), key="x")],
)
assert len(f1.edges) == 1
assert len(f1.tasks) == 2
def test_viz_reflects_multiple_mapping_if_flow_state_provided(self):
ipython = MagicMock(
get_ipython=lambda: MagicMock(config=dict(IPKernelApp=True)))
add = AddTask(name="a_nice_task")
list_task = Task(name="a_list_task")
map_state1 = Mapped(map_states=[Success(), TriggerFailed()])
map_state2 = Mapped(map_states=[Success(), Failed()])
with patch.dict("sys.modules", IPython=ipython):
with Flow(name="test") as f:
first_res = add.map(x=list_task, y=8)
with pytest.warns(
UserWarning
): # making a copy of a task with dependencies
res = first_res.map(x=first_res, y=9)
graph = f.visualize(flow_state=Success(
result={
res: map_state1,
list_task: Success(),
first_res: map_state2,
}))
assert "{first} -> {second} [label=x style=dashed]".format(
first=str(id(first_res)) + "0", second=str(id(res)) + "0")
assert "{first} -> {second} [label=x style=dashed]".format(
first=str(id(first_res)) + "1", second=str(id(res)) + "1")
def test_equality_based_on_reference_tasks(self):
f1 = Flow(name="test")
f2 = Flow(name="test")
t1 = Task()
t2 = Task()
t3 = Task()
for f in [f1, f2]:
f.add_edge(t1, t2)
f.add_edge(t1, t3)
f1.set_reference_tasks([t2])
assert f1 != f2
f2.set_reference_tasks([t2])
assert f1 == f2
def test_setting_reference_tasks_clears_cache(self):
f = Flow(name="test")
t1 = Task()
f.add_task(t1)
f._cache[1] = 2
f.set_reference_tasks([t1])
assert 1 not in f._cache
def test_does_not_look_for_file_when_no_result_handler_given(self):
task = Task(name="Task")
task_runner = DSTaskRunner(task) # Does not get upstream_states by
# default, so normally would throw an error if a result handler was passed
old_state = Pending()
new_state = Running()
dsh.checkpoint_handler(task_runner, old_state, new_state)
def test_replace_replaces_all_the_things(self):
with Flow(name="test") as f:
t1 = Task(name="t1")()
t2 = Task(name="t2")(upstream_tasks=[t1])
t3 = Task(name="t3")
f.set_reference_tasks([t1])
f.replace(t1, t3)
assert f.tasks == {t2, t3}
assert {e.upstream_task for e in f.edges} == {t3}
assert {e.downstream_task for e in f.edges} == {t2}
assert f.reference_tasks() == {t3}
assert f.terminal_tasks() == {t2}
with pytest.raises(ValueError):
f.edges_to(t1)
def test_viz_reflects_mapping_if_flow_state_provided(self):
ipython = MagicMock(
get_ipython=lambda: MagicMock(config=dict(IPKernelApp=True)))
add = AddTask(name="a_nice_task")
list_task = Task(name="a_list_task")
map_state = Mapped(map_states=[Success(), Failed()])
with patch.dict("sys.modules", IPython=ipython):
with Flow(name="test") as f:
res = add.map(x=list_task, y=8)
graph = f.visualize(flow_state=Success(result={
res: map_state,
list_task: Success()
}))
# one colored node for each mapped result
assert 'label="a_nice_task <map>" color="#00800080"' in graph.source
assert 'label="a_nice_task <map>" color="#FF000080"' in graph.source
assert 'label=a_list_task color="#00800080"' in graph.source
assert 'label=8 color="#00000080"' in graph.source
# two edges for each input to add()
for var in ["x", "y"]:
for index in [0, 1]:
assert "{0} [label={1} style=dashed]".format(
index, var) in graph.source
def test_sorted_tasks_with_start_task():
"""
t1 -> t2 -> t3 -> t4
t3 -> t5
"""
f = Flow(name="test")
t1 = Task("1")
t2 = Task("2")
t3 = Task("3")
t4 = Task("4")
t5 = Task("5")
f.add_edge(t1, t2)
f.add_edge(t2, t3)
f.add_edge(t3, t4)
f.add_edge(t3, t5)
assert set(f.sorted_tasks(root_tasks=[])) == set([t1, t2, t3, t4, t5])
assert set(f.sorted_tasks(root_tasks=[t3])) == set([t3, t4, t5])
def test_copy():
with Flow(name="test") as f:
t1 = Task()
t2 = Task()
t3 = Task()
f.add_edge(t1, t2)
f.add_edge(t2, t3)
f.set_reference_tasks([t1])
f2 = f.copy()
assert f2 == f
f.add_edge(Task(), Task())
assert len(f2.tasks) == len(f.tasks) - 2
assert len(f2.edges) == len(f.edges) - 1
assert f.reference_tasks() == f2.reference_tasks() == set([t1])
def test_set_dependencies_adds_all_arguments_to_flow():
f = Flow(name="test")
class ArgTask(Task):
def run(self, x):
return x
t1 = ArgTask()
t2 = Task()
t3 = Task()
t4 = Task()
f.set_dependencies(
task=t1, upstream_tasks=[t2], downstream_tasks=[t3], keyword_tasks={"x": t4}
)
assert f.tasks == set([t1, t2, t3, t4])
def test_raises_appropriate_error_when_incompatible_handler_given(self):
task = Task(name="Task", result_handler=LocalResultHandler())
task_runner = DSTaskRunner(task)
task_runner.upstream_states = {}
old_state = Pending()
new_state = Running()
with pytest.raises(TypeError):
dsh.checkpoint_handler(task_runner, old_state, new_state)
def test_sorted_tasks_with_ambiguous_sort():
"""
t1 -> bottleneck
t2 -> bottleneck
t3 -> bottleneck
bottleneck -> t4
bottleneck -> t5
bottleneck -> t6
"""
f = Flow(name="test")
t1 = Task("1")
t2 = Task("2")
t3 = Task("3")
t4 = Task("4")
t5 = Task("5")
t6 = Task("6")
bottleneck = Task("bottleneck")
f.add_edge(t1, bottleneck)
f.add_edge(t2, bottleneck)
f.add_edge(t3, bottleneck)
f.add_edge(bottleneck, t4)
f.add_edge(bottleneck, t5)
f.add_edge(bottleneck, t6)
tasks = f.sorted_tasks()
assert set(tasks[:3]) == set([t1, t2, t3])
assert list(tasks)[3] is bottleneck
assert set(tasks[4:]) == set([t4, t5, t6])
def test_add_edge_raise_error_for_downstream_parameter():
f = Flow(name="test")
t = Task()
p = Parameter("p")
with pytest.raises(ValueError) as exc:
f.add_edge(upstream_task=t, downstream_task=p)
assert "can not have upstream dependencies" in str(exc.value)
def test_cache_sorted_tasks(self):
f = Flow(name="test")
t1 = Task()
t2 = Task()
t3 = Task()
f.add_edge(t1, t2)
f.sorted_tasks()
# check that cache holds result
key = ("_sorted_tasks", (("root_tasks", ()), ))
assert f._cache[key] == (t1, t2)
# check that cache is read
f._cache[key] = 1
assert f.sorted_tasks() == 1
f.add_edge(t2, t3)
assert f.sorted_tasks() == (t1, t2, t3)
def test_errors_when_regular_runner_is_used(self):
task = Task(name="Task",
result_handler=PandasResultHandler("dummy.csv", "csv"))
task_runner = TaskRunner(task)
old_state = Pending()
new_state = Running()
with pytest.raises(TypeError):
dsh.checkpoint_handler(task_runner, old_state, new_state)
def test_bad_flow_runner_code_still_returns_state_obj():
class BadFlowRunner(prefect.engine.flow_runner.FlowRunner):
def initialize_run(self, *args, **kwargs):
import blig # will raise ImportError
f = Flow(name="test", tasks=[Task()])
res = f.run(runner_cls=BadFlowRunner)
assert isinstance(res, State)
assert res.is_failed()
assert isinstance(res.result, ImportError)
def test_add_edge_raise_error_for_duplicate_key_if_validate():
f = Flow(name="test")
t = Task()
a = AddTask()
f.add_edge(upstream_task=t, downstream_task=a, key="x")
with pytest.raises(ValueError) as exc:
f.add_edge(upstream_task=t, downstream_task=a, key="x", validate=True)
assert "already been assigned" in str(exc.value)
def test_cache_task_ids(self):
f = Flow(name="test")
t1 = Task()
t2 = Task()
t3 = Task()
f.add_edge(t1, t2)
ids = f.task_ids
# check that cache holds result
key = ("task_ids", ())
assert f._cache[key] == ids
# check that cache is read
f._cache[key] = 1
assert f.task_ids == 1
f.add_edge(t2, t3)
assert len(f.task_ids) == 3
def test_cache_terminal_tasks(self):
f = Flow(name="test")
t1 = Task()
t2 = Task()
t3 = Task()
f.add_edge(t1, t2)
f.terminal_tasks()
# check that cache holds result
key = ("terminal_tasks", ())
assert f._cache[key] == set([t2])
# check that cache is read
f._cache[key] = 1
assert f.terminal_tasks() == 1
f.add_edge(t2, t3)
assert f.terminal_tasks() == set([t3])
def test_chain_works_in_flow_context_without_duplication():
@task
def do_nothing():
pass
with Flow(name="test") as f:
f.chain(do_nothing(), do_nothing(), do_nothing(), Task())
assert len(f.tasks) == 4
assert len(f.edges) == 3
def test_visualize_raises_informative_importerror_without_graphviz(
self, monkeypatch):
f = Flow(name="test")
f.add_task(Task())
with monkeypatch.context() as m:
m.setattr(sys, "path", "")
with pytest.raises(ImportError) as exc:
f.visualize()
assert "pip install prefect[viz]" in repr(exc.value)
def test_viz_returns_graph_object_if_in_ipython(self):
import graphviz
ipython = MagicMock(
get_ipython=lambda: MagicMock(config=dict(IPKernelApp=True)))
with patch.dict("sys.modules", IPython=ipython):
f = Flow(name="test")
f.add_task(Task(name="a_nice_task"))
graph = f.visualize()
assert "label=a_nice_task" in graph.source
assert "shape=ellipse" in graph.source
