def test_skip_validation_in_init_with_kwarg():
t1, t2 = Task(), Task() # these tasks don't support keyed edges
e1, e2 = Edge(t1, t2), Edge(t2, t1)
with pytest.raises(ValueError):
Flow(name="test", edges=[e1, e2], validate=True)
assert Flow(name="test", edges=[e1, e2], validate=False)
def test_works_with_multiple_upstream_states(self):
upstream_task_1 = Task(name="upstream_task_one")
upstream_state_1 = State(result=1)
upstream_task_2 = Task(name="upstream_task_two")
upstream_state_2 = State(result=2)
downstream_task = Task(name="downstream_task")
upstream_states = {
Edge(upstream_task_1, downstream_task, key="var_1"):
upstream_state_1,
Edge(upstream_task_2, downstream_task, key="var_2"):
upstream_state_2
}
mapping = dsh._create_input_mapping(upstream_states)
assert mapping == {"var_1": 1, "var_2": 2}
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_add_edge_returns_edge():
f = Flow(name="test")
t1 = Task()
t2 = Task()
edge = Edge(t1, t2)
added_edge = f.add_edge(upstream_task=t1, downstream_task=t2)
assert edge == added_edge
assert added_edge in f.edges
assert edge in f.edges
def test_binding_a_task_with_var_kwargs_expands_the_kwargs():
class KwargsTask(Task):
def run(self, **kwargs):
return kwargs
t1 = Task()
t2 = Task()
t3 = Task()
kw = KwargsTask()
with Flow(name="test") as f:
kw.bind(a=t1, b=t2, c=t3)
assert t1 in f.tasks
assert t2 in f.tasks
assert t3 in f.tasks
assert Edge(t1, kw, key="a") in f.edges
assert Edge(t2, kw, key="b") in f.edges
assert Edge(t3, kw, key="c") in f.edges
def add_edge(
self,
upstream_task: Task,
downstream_task: Task,
key: str = None,
mapped: bool = False,
validate: bool = None,
) -> Edge:
"""
Add an edge in the flow between two tasks. All edges are directed beginning with
an upstream task and ending with a downstream task.
Args:
- upstream_task (Task): The task that the edge should start from
- downstream_task (Task): The task that the edge should end with
- key (str, optional): The key to be set for the new edge; the result of the upstream task
will be passed to the downstream task's `run()` method under this keyword argument
- mapped (bool, optional): Whether this edge represents a call to `Task.map()`; defaults to `False`
- validate (bool, optional): Whether or not to check the validity of
the flow (e.g., presence of cycles and illegal keys). Defaults to the value
of `eager_edge_validation` in your prefect configuration file.
Returns:
- prefect.core.edge.Edge: The `Edge` object that was successfully added to the flow
Raises:
- ValueError: if the `downstream_task` is of type `Parameter`
- ValueError: if the edge exists with this `key` and `downstream_task`
"""
if validate is None:
validate = cast(bool, prefect.config.flows.eager_edge_validation)
if isinstance(downstream_task, Parameter):
raise ValueError(
"Parameters must be root tasks and can not have upstream dependencies."
)
self.add_task(upstream_task)
self.add_task(downstream_task)
# we can only check the downstream task's edges once it has been added to the
# flow, so we need to perform this check here and not earlier.
if validate and key and key in {
e.key
for e in self.edges_to(downstream_task)
}:
raise ValueError(
'Argument "{a}" for task {t} has already been assigned in '
"this flow. If you are trying to call the task again with "
"new arguments, call Task.copy() before adding the result "
"to this flow.".format(a=key, t=downstream_task))
edge = Edge(
upstream_task=upstream_task,
downstream_task=downstream_task,
key=key,
mapped=mapped,
)
self.edges.add(edge)
# check that the edges are valid keywords by binding them
if validate and key is not None:
edge_keys = {
e.key: None
for e in self.edges_to(downstream_task) if e.key is not None
}
inspect.signature(downstream_task.run).bind_partial(**edge_keys)
self._cache.clear()
# check for cycles
if validate:
self.validate()
return edge
flow.add_edge(upstream_task=e, downstream_task=t, key="x") flow.add_edge(upstream_task=t, downstream_task=l, key="x") state = flow.run() # Testing state assert state.is_successful assert state.result[e].is_successful assert state.result[t].is_successful assert state.result[l].is_successful # Testing results assert state.result[e].result == 10 assert state.result[t].result == 20 assert state.result[l].result == None # Flow composition assert e in flow.tasks assert t in flow.tasks assert l in flow.tasks assert len(flow.tasks) == 3 assert flow.root_tasks() == set([e]) assert flow.terminal_tasks() == set([l]) assert len(flow.edges) == 2 assert Edge(upstream_task=e, downstream_task=t, key="x") in flow.edges assert Edge(upstream_task=t, downstream_task=l, key="x") in flow.edges