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