Пример #1
0
    def test_as_task_toggles_constants(self):
        with Flow("test") as f:
            t = tasks.as_task(4)
            s = tasks.as_task(5, convert_constants=False)

        assert isinstance(t, Task)
        assert t.name == "4"

        assert not isinstance(s, Task)
        assert s == 5
Пример #2
0
def test_list_of_tasks():
    """
    Test that a list of tasks can be set as a switch condition
    """

    with Flow(name="test") as flow:
        condition = Condition()
        true_branch = [SuccessTask(), SuccessTask()]
        list_task = as_task(true_branch)
        false_branch = SuccessTask()
        ifelse(condition, list_task, false_branch)

    with prefect.context(CONDITION=True):
        state = flow.run()

        for t in true_branch:
            assert isinstance(state.result[t], Success)
        assert isinstance(state.result[false_branch], Skipped)

    with prefect.context(CONDITION=False):
        state = flow.run()

        for t in true_branch:
            # the tasks in the list ran becuase they have no upstream dependencies.
            assert isinstance(state.result[t], Success)
        # but the list itself skipped
        assert isinstance(state.result[list_task], Skipped)
        assert isinstance(state.result[false_branch], Success)
Пример #3
0
 def test_ordered_collections(self):
     """
     Tests that ordered collections maintain order
     """
     val = [[list(range(100))]]
     with Flow("test") as f:
         task = tasks.as_task(val)
         f.add_task(task)
     assert f.run().result[task].result == val
Пример #4
0
    def replace(self, old: Task, new: Task, validate: bool = True) -> None:
        """
        Performs an inplace replacement of the old task with the provided new task.

        Args:
            - old (Task): the old task to replace
            - new (Task): the new task to replace the old with; if not a Prefect
                Task, Prefect will attempt to convert it to one
            - validate (boolean, optional): whether to validate the Flow after
                the replace has been completed; defaults to `True`

        Raises:
            - ValueError: if the `old` task is not a part of this flow
        """
        if old not in self.tasks:
            raise ValueError("Task {t} was not found in Flow {f}".format(
                t=old, f=self))

        new = as_task(new)

        # update tasks
        self.tasks.remove(old)
        self.add_task(new)

        self._cache.clear()

        affected_edges = {e for e in self.edges if old in e.tasks}

        # remove old edges
        for edge in affected_edges:
            self.edges.remove(edge)

        # replace with new edges
        for edge in affected_edges:
            upstream = new if edge.upstream_task == old else edge.upstream_task
            downstream = new if edge.downstream_task == old else edge.downstream_task
            self.add_edge(
                upstream_task=upstream,
                downstream_task=downstream,
                key=edge.key,
                mapped=edge.mapped,
                validate=False,
            )

        # update auxiliary task collections
        ref_tasks = self.reference_tasks()
        new_refs = [t for t in ref_tasks if t != old
                    ] + ([new] if old in ref_tasks else [])
        self.set_reference_tasks(new_refs)

        if validate:
            self.validate()
Пример #5
0
    def test_apply_map_simple(self, api):
        if api == "functional":

            def func(x, y, z):
                a = add(x, y)
                a.name = "add-a"
                b = add(a, z)
                b.name = "add-b"
                c = add(b, 1)
                c.name = "add-c"
                return inc(c)

            with Flow("test") as flow:
                y = ranged(3)
                z = edges.unmapped(1)
                res = apply_map(func, range(3, 6), y=y, z=z)
        else:

            def func(x, y, z, flow):
                a = add.copy(name="add-a").bind(x, y, flow=flow)
                b = add.copy(name="add-b").bind(a, z, flow=flow)
                c = add.copy(name="add-c").bind(b, 1, flow=flow)
                return inc.copy().bind(c, flow=flow)

            flow = Flow("test")
            y = ranged.copy().bind(3, flow=flow)
            z = edges.unmapped(tasks.as_task(1, flow=flow))
            res = apply_map(func, range(3, 6), y=y, z=z, flow=flow)

        consts = {t.name: c for t, c in flow.constants.items()}
        assert consts == {
            "ranged": {
                "n": 3
            },
            "add-b": {
                "y": 1
            },
            "add-c": {
                "y": 1
            }
        }

        for task in flow.tasks:
            if task.name != "ranged":
                for e in flow.edges_to(task):
                    assert e.mapped

        state = flow.run()
        assert state.result[res].result == [6, 8, 10]
Пример #6
0
    def test_as_task_with_basic_python_objs(self, obj):
        @tasks.task
        def return_val(x):
            "Necessary because constant tasks aren't tracked inside the flow"
            return x

        with Flow("test") as f:
            t = tasks.as_task(obj)
            val = return_val(t)

        assert isinstance(t, Task)
        res = FlowRunner(f).run(return_tasks=[val])

        assert res.is_successful()
        assert res.result[val].result == obj
Пример #7
0
 def test_nested_collections(self, val):
     with Flow("test") as f:
         task = tasks.as_task(val)
         f.add_task(task)
     assert f.run().result[task].result == val
Пример #8
0
 def test_nested_collections_of_mixed_constants_are_not_constants(
         self, val):
     with Flow("test"):
         task = tasks.as_task(val)
     assert not isinstance(task, Constant)
Пример #9
0
 def test_nested_collections_of_constants_are_constants(self, val):
     task = tasks.as_task(val)
     assert isinstance(task, Constant)
     assert task.value == val
Пример #10
0
 def test_as_task_doesnt_label_tasks_as_auto_generated(self):
     t = Task()
     assert t.auto_generated is False
     assert tasks.as_task(t).auto_generated is False
Пример #11
0
    def test_as_task_toggles_constants(self):
        with Flow("test"):
            t = tasks.as_task(4)

        assert isinstance(t, Task)
        assert t.name == "4"
Пример #12
0
    def set_dependencies(
        self,
        task: object,
        upstream_tasks: Iterable[object] = None,
        downstream_tasks: Iterable[object] = None,
        keyword_tasks: Mapping[str, object] = None,
        mapped: bool = False,
        validate: bool = None,
    ) -> None:
        """
        Convenience function for adding task dependencies.

        Args:
            - task (object): a Task that will become part of the Flow. If the task is not a
                Task subclass, Prefect will attempt to convert it to one.
            - upstream_tasks ([object], optional): Tasks that will run before the task runs. If any task
                is not a Task subclass, Prefect will attempt to convert it to one.
            - downstream_tasks ([object], optional): Tasks that will run after the task runs. If any task
                is not a Task subclass, Prefect will attempt to convert it to one.
            - keyword_tasks ({key: object}, optional): The results of these tasks
                will be provided to the task under the specified keyword
                arguments. If any task is not a Task subclass, Prefect will attempt to
                convert it to one.
            - mapped (bool, optional): Whether the upstream tasks (both keyed
                and non-keyed) should be mapped over; defaults to `False`. If `True`, any
                tasks wrapped in the `prefect.utilities.tasks.unmapped` container will
                _not_ be mapped over.
            - validate (bool, optional): Whether or not to check the validity of
                the flow (e.g., presence of cycles).  Defaults to the value of `eager_edge_validation`
                in your Prefect configuration file.

        Returns:
            - None
        """

        task = as_task(task)
        assert isinstance(task, Task)  # mypy assert

        # add the main task (in case it was called with no arguments)
        self.add_task(task)

        # add upstream tasks
        for t in upstream_tasks or []:
            is_mapped = mapped & (not isinstance(t, unmapped))
            t = as_task(t)
            assert isinstance(t, Task)  # mypy assert
            self.add_edge(
                upstream_task=t,
                downstream_task=task,
                validate=validate,
                mapped=is_mapped,
            )

        # add downstream tasks
        for t in downstream_tasks or []:
            t = as_task(t)
            assert isinstance(t, Task)  # mypy assert
            self.add_edge(upstream_task=task,
                          downstream_task=t,
                          validate=validate)

        # add data edges to upstream tasks
        for key, t in (keyword_tasks or {}).items():
            is_mapped = mapped & (not isinstance(t, unmapped))
            t = as_task(t)
            assert isinstance(t, Task)  # mypy assert
            self.add_edge(
                upstream_task=t,
                downstream_task=task,
                key=key,
                validate=validate,
                mapped=is_mapped,
            )
Пример #13
0
 def test_as_task_unpacks_unmapped_objects(self):
     t1 = Task()
     unmapped_t1 = tasks.unmapped(t1)
     assert tasks.as_task(t1) is t1
     assert tasks.as_task(t1).auto_generated is False
Пример #14
0
 def test_as_task_unpacks_unmapped_objects(self, edge_annotation):
     t1 = Task()
     annotated = edge_annotation(t1)
     assert tasks.as_task(annotated) is t1
     assert tasks.as_task(annotated).auto_generated is False
Пример #15
0
 def test_as_task_unpacks_unmapped_objects(self):
     t1 = Task()
     unmapped_t1 = tasks.unmapped(t1)
     assert tasks.as_task(t1) is t1