test_task.py

import inspect
import json
import logging
import uuid
from datetime import timedelta
from typing import Any, Union
from unittest.mock import MagicMock

import pytest

import prefect
from prefect.core import Edge, Flow, Parameter, Task
from prefect.engine.cache_validators import all_inputs, duration_only, never_use
from prefect.engine.result_handlers import JSONResultHandler, ResultHandler
from prefect.utilities.configuration import set_temporary_config
from prefect.utilities.tasks import task


class AddTask(Task):
    def run(self, x, y=1):
        return x + y


class TestCreateTask:
    """Test various Task constructors"""

    def test_create_task_with_no_args(self):
        """Tasks can be created with no arguments"""
        assert Task()

    def test_create_task_with_name(self):
        t1 = Task()
        assert t1.name == "Task"

        t2 = Task(name="test")
        assert t2.name == "test"

    def test_create_task_with_slug(self):
        t1 = Task()
        assert t1.slug
        assert uuid.UUID(t1.slug)  # slug is a UUID

        t2 = Task(slug="test")
        assert t2.slug == "test"

    def test_create_task_with_max_retries(self):
        t1 = Task()
        assert t1.max_retries == 0

        t2 = Task(max_retries=5, retry_delay=timedelta(0))
        assert t2.max_retries == 5

    def test_create_task_with_retry_delay(self):
        t2 = Task(retry_delay=timedelta(seconds=30))
        assert t2.retry_delay == timedelta(seconds=30)

    def test_create_task_with_max_retries_and_no_retry_delay(self):
        with pytest.raises(ValueError):
            Task(max_retries=1, retry_delay=None)

    def test_create_task_with_timeout(self):
        t1 = Task()
        assert t1.timeout == None

        with pytest.raises(TypeError):
            Task(timeout=0.5)

        t3 = Task(timeout=1)
        assert t3.timeout == 1

        with set_temporary_config({"tasks.defaults.timeout": 3}):
            t4 = Task()
            assert t4.timeout == 3

    def test_create_task_with_trigger(self):
        t1 = Task()
        assert t1.trigger is prefect.triggers.all_successful

        t2 = Task(trigger=prefect.triggers.all_failed)
        assert t2.trigger == prefect.triggers.all_failed

    def test_create_task_without_state_handler(self):
        assert Task().state_handlers == []

    @pytest.mark.parametrize("handlers", [[lambda *a: 1], [lambda *a: 1, lambda *a: 2]])
    def test_create_task_with_state_handler(self, handlers):
        assert Task(state_handlers=handlers).state_handlers == handlers

    def test_create_task_with_on_failure(self):
        t = Task(on_failure=lambda *args: None)
        assert len(t.state_handlers) == 1

    def test_create_task_illegal_handler(self):
        with pytest.raises(TypeError):
            Task(state_handlers=lambda *a: 1)

    def test_class_creation_rejects_reserved_init_kwargs(self):
        with pytest.raises(ValueError):

            class BadTask(Task):
                def __init__(self, name):
                    pass

        with pytest.raises(ValueError):

            class BadTask(Task):
                def __init__(self, checkpoint):
                    pass

    @pytest.mark.parametrize(
        "arg", [a for a in inspect.getfullargspec(Task.__init__).args if a != "self"]
    )
    def test_validator_validates_all_task_init_kwargs(self, arg, monkeypatch):
        monkeypatch.setattr(
            "prefect.core.task.inspect.getfullargspec",
            MagicMock(return_value=MagicMock(args=[arg])),
        )
        with pytest.raises(ValueError) as exc:
            prefect.core.task._validate_init_signature(lambda: None)
        assert arg in str(exc.value)

    def test_class_instantiation_rejects_varargs(self):
        with pytest.raises(ValueError):

            class VarArgsTask(Task):
                def run(self, x, *y):
                    pass

    def test_class_instantiation_rejects_mapped_kwarg(self):
        with pytest.raises(ValueError):

            class MappedTasks(Task):
                def run(self, x, mapped):
                    pass

        with pytest.raises(ValueError):

            class MappedTasks(Task):
                def run(self, x, mapped=None):
                    pass

    def test_class_instantiation_rejects_mapped_kwarg_decorator(self):
        with pytest.raises(ValueError):

            @task
            def run(x, mapped):
                pass

        with pytest.raises(ValueError):

            @task
            def run(x, mapped=None):
                pass

    def test_class_instantiation_rejects_upstream_tasks_kwarg(self):
        with pytest.raises(ValueError):

            class UpstreamTasks(Task):
                def run(self, x, upstream_tasks):
                    pass

        with pytest.raises(ValueError):

            class UpstreamTasks(Task):
                def run(self, x, upstream_tasks=None):
                    pass

    def test_class_instantiation_rejects_upstream_tasks_kwarg_decorator(self):
        with pytest.raises(ValueError):

            @task
            def run(x, upstream_tasks):
                pass

        with pytest.raises(ValueError):

            @task
            def run(x, upstream_tasks=None):
                pass

    def test_class_instantiation_rejects_flow_kwarg(self):
        with pytest.raises(ValueError):

            class FlowTasks(Task):
                def run(self, x, flow):
                    pass

        with pytest.raises(ValueError):

            class FlowTasks(Task):
                def run(self, x, flow=None):
                    pass

    def test_class_instantiation_rejects_flow_kwarg_decorator(self):
        with pytest.raises(ValueError):

            @task
            def run(x, flow):
                pass

        with pytest.raises(ValueError):

            @task
            def run(x, flow=None):
                pass

    def test_class_instantiation_rejects_task_args_kwarg(self):
        with pytest.raises(ValueError):

            class TaskArgs(Task):
                def run(self, x, task_args):
                    pass

        with pytest.raises(ValueError):

            class TaskArgs(Task):
                def run(self, x, task_args=None):
                    pass

    def test_class_instantiation_rejects_task_args_kwarg_decorator(self):
        with pytest.raises(ValueError):

            @task
            def run(x, task_args):
                pass

        with pytest.raises(ValueError):

            @task
            def run(x, task_args=None):
                pass

    def test_create_task_with_and_without_cache_for(self):
        t1 = Task()
        assert t1.cache_validator is never_use
        t2 = Task(cache_for=timedelta(days=1))
        assert t2.cache_validator is duration_only
        t3 = Task(cache_for=timedelta(days=1), cache_validator=all_inputs)
        assert t3.cache_validator is all_inputs

    def test_bad_cache_kwarg_combo(self):
        with pytest.warns(UserWarning, match=".*Task will not be cached.*"):
            Task(cache_validator=all_inputs)

    def test_create_task_with_and_without_result_handler(self):
        t1 = Task()
        assert t1.result_handler is None
        t2 = Task(result_handler=JSONResultHandler())
        assert isinstance(t2.result_handler, ResultHandler)
        assert isinstance(t2.result_handler, JSONResultHandler)

    def test_create_parameter_uses_json_result_handler(self):
        p = Parameter("p")
        assert isinstance(p.result_handler, JSONResultHandler)

    def test_create_task_with_and_without_checkpoint(self):
        t = Task()
        assert t.checkpoint is False

        s = Task(checkpoint=True)
        assert s.checkpoint is True

        with set_temporary_config({"tasks.defaults.checkpoint": True}):
            r = Task()
        assert r.checkpoint is True

    @pytest.mark.xfail(reason="UX improvement for Core")
    def test_create_parameter_always_checkpoints(self):
        with set_temporary_config({"tasks.defaults.checkpoint": False}):
            p = Parameter("p")
        assert p.checkpoint is True


def test_task_has_logger():
    t = Task()
    assert isinstance(t.logger, logging.Logger)
    assert t.logger.name == "prefect.Task"


def test_task_produces_no_result():
    t = Task()
    assert t.run() is None


def test_task_is_not_iterable():
    t = Task()
    with pytest.raises(TypeError):
        list(t)


def test_tags_are_added_when_arguments_are_bound():
    t1 = AddTask(tags=["math"])
    t2 = AddTask(tags=["math"])

    with prefect.context(tags=["test"]):
        with Flow(name="test"):
            t1.bind(1, 2)
            t3 = t2(1, 2)

    assert t1.tags == {"math", "test"}
    assert t3.tags == {"math", "test"}


def test_tags():
    t1 = Task()
    assert t1.tags == set()

    with pytest.raises(TypeError):
        Task(tags="test")

    t3 = Task(tags=["test", "test2", "test"])
    assert t3.tags == set(["test", "test2"])

    with prefect.context(tags=["test"]):
        t4 = Task()
        assert t4.tags == set(["test"])

    with prefect.context(tags=["test1", "test2"]):
        t5 = Task(tags=["test3"])
        assert t5.tags == set(["test1", "test2", "test3"])


class TestInputsOutputs:
    class add(Task):
        def run(self, x, y: int = 1) -> int:
            return x + y

    @task
    def mult(x, y: int = 1) -> int:
        return x * y

    def test_inputs(self):
        assert self.add().inputs() == dict(
            x=dict(type=Any, required=True, default=None),
            y=dict(type=int, required=False, default=1),
        )

    def test_inputs_task_decorator(self):
        with Flow("test"):
            assert self.mult(x=1).inputs() == dict(
                x=dict(type=Any, required=True, default=None),
                y=dict(type=int, required=False, default=1),
            )

    def test_outputs(self):
        assert self.add().outputs() == int

    def test_outputs_task_decorator(self):
        with Flow("test"):
            assert self.mult(x=1).outputs() == int


class TestTaskCopy:
    def test_copy_copies(self):
        class CopyTask(Task):
            class_attr = 42

            def __init__(self, instance_val, **kwargs):
                self.instance_val = instance_val
                super().__init__(**kwargs)

            def run(self, run_val):
                return (run_val, self.class_attr, self.instance_val)

        ct = CopyTask("username")
        other = ct.copy()
        assert isinstance(other, CopyTask)
        assert ct is not other
        assert hash(ct) != hash(other)
        assert ct != other
        assert other.run("pass") == ("pass", 42, "username")

    def test_copy_warns_if_dependencies_in_active_flow(self):
        t1 = Task()
        t2 = Task()

        with Flow(name="test"):
            t1.set_dependencies(downstream_tasks=[t2])
            with pytest.warns(UserWarning):
                t1.copy()

            with Flow(name="test"):
                # no dependencies in this flow
                t1.copy()

    def test_copy_changes_slug(self):
        t1 = Task(slug="test")
        t2 = t1.copy()
        assert t1.slug == "test"
        assert t1.slug != t2.slug

    def test_copy_accepts_task_args(self):
        t = Task()
        t2 = t.copy(name="new-task")
        t3 = t.copy(**{"max_retries": 4200})

        assert t2.name == "new-task"
        assert t3.max_retries == 4200

    def test_copy_accepts_slug_as_task_args(self):
        t = Task(slug="test")
        t2 = t.copy(slug="test-2")
        assert t.slug == "test"
        assert t2.slug == "test-2"


def test_task_has_slug():
    t1 = Task()
    t2 = Task()

    assert t1.slug and t1.slug != t2.slug


class TestDependencies:
    def test_set_downstream(self):
        f = Flow(name="test")
        t1 = Task()
        t2 = Task()
        t1.set_downstream(t2, flow=f)
        assert Edge(t1, t2) in f.edges

    def test_set_downstream_context(self):
        with Flow(name="test") as f:
            t1 = Task()
            t2 = Task()
            t1.set_downstream(t2)
            assert Edge(t1, t2) in f.edges

    def test_set_downstream_no_flow(self):
        f = Flow(name="test")
        t1 = Task()
        t2 = Task()
        with pytest.raises(ValueError) as exc:
            t1.set_downstream(t2)
        assert "No Flow was passed" in str(exc.value)

    @pytest.mark.parametrize(
        "props", [{"mapped": True}, {"key": "x"}, {"key": "x", "mapped": True}]
    )
    def test_set_downstream_with_properties(self, props):
        with Flow(name="test") as f:
            t1 = Task()
            t2 = Task()
            t1.set_downstream(t2, **props)
            assert Edge(t1, t2, **props) in f.edges

    def test_set_upstream(self):
        f = Flow(name="test")
        t1 = Task()
        t2 = Task()
        t2.set_upstream(t1, flow=f)
        assert Edge(t1, t2) in f.edges

    def test_set_upstream_context(self):
        with Flow(name="test") as f:
            t1 = Task()
            t2 = Task()
            t2.set_upstream(t1)
            assert Edge(t1, t2) in f.edges

    def test_set_upstream_no_flow(self):
        f = Flow(name="test")
        t1 = Task()
        t2 = Task()
        with pytest.raises(ValueError) as exc:
            t2.set_upstream(t1)
        assert "No Flow was passed" in str(exc.value)

    @pytest.mark.parametrize(
        "props", [{"mapped": True}, {"key": "x"}, {"key": "x", "mapped": True}]
    )
    def test_set_upstream_with_properties(self, props):
        with Flow(name="test") as f:
            t1 = Task()
            t2 = Task()
            t2.set_upstream(t1, **props)
            assert Edge(t1, t2, **props) in f.edges


class TestSerialization:
    def test_serialization(self):
        t = Task(name="test")
        s = t.serialize()

        assert isinstance(s, dict)
        assert s["slug"] == t.slug
        assert s["type"] == "prefect.core.task.Task"
        assert s["name"] == t.name

    def test_subclass_serialization(self):
        class NewTask(Task):
            pass

        s = NewTask().serialize()

        assert isinstance(s, dict)
        assert s["type"].endswith(".NewTask")

    def test_deserialization(self):
        t = Task(name="test")
        s = t.serialize()
        t2 = prefect.serialization.task.TaskSchema().load(s)
        assert isinstance(t2, Task)
        assert t2.name == t.name

    def test_subclass_deserialization(self):
        class NewTask(Task):
            pass

        t = NewTask(name="test")
        s = t.serialize()
        t2 = prefect.serialization.task.TaskSchema().load(s)
        assert type(t2) is Task
        assert not isinstance(t2, NewTask)
        assert t2.name == t.name

    def test_parameter_serialization(self):
        p = Parameter(name="p")
        serialized = p.serialize()
        assert serialized["name"] == "p"
        assert serialized["default"] is None
        assert serialized["required"] is True

    def test_parameter_deserialization(self):
        p = Parameter(name="p")
        serialized = p.serialize()
        p2 = prefect.serialization.task.ParameterSchema().load(serialized)
        assert isinstance(p2, Parameter)
        assert p2.name == p.name
        assert p2.required == p.required
        assert p2.default == p.default


class TestTaskArgs:
    def test_task_args_raises_for_non_attrs(self):
        t = Task()
        with Flow(name="test") as f:
            with pytest.raises(AttributeError) as exc:
                res = t(task_args={"foo": "bar"})

        assert "foo" in str(exc.value)

    @pytest.mark.parametrize(
        "attr,val",
        [
            ("name", "foo-bar"),
            ("slug", "foo-bar"),
            ("max_retries", 4200),
            ("retry_delay", timedelta(seconds=1)),
            ("timeout", 12),
            ("skip_on_upstream_skip", False),
            ("cache_for", timedelta(seconds=1)),
        ],
    )
    def test_task_args_sets_new_attrs(self, attr, val):
        t = Task()
        with Flow(name="test") as f:
            res = t(task_args={attr: val})

        assert getattr(f.tasks.pop(), attr) == val

    def test_tags_are_appended_to_when_updating_with_task_args(self):
        t = AddTask(tags=["math"])

        with prefect.context(tags=["test"]):
            with Flow(name="test"):
                t2 = t(1, 2, task_args={"name": "test-tags", "tags": ["new-tag"]})

        assert t2.tags == {"math", "test", "new-tag"}