Repository: PrefectHQ/prefect

Issue:
Allow SUCCESS signals to carry results
Currently, tasks can `raise signals.SUCCESS()` but not provide a result. It would be great if a result could be passed and applied to the resulting state.


Discussion:

Resolving patch:
diff --git a/src/prefect/engine/signals.py b/src/prefect/engine/signals.py
--- a/src/prefect/engine/signals.py
+++ b/src/prefect/engine/signals.py
@@ -19,9 +19,9 @@ class PrefectStateSignal(PrefectError):
 
     _state_cls = state.State
 
-    def __init__(self, message=None, **kwargs) -> None:  # type: ignore
-        super().__init__(message, **kwargs)  # type: ignore
-        self.state = self._state_cls(message=self)
+    def __init__(self, message=None, *args, **kwargs) -> None:  # type: ignore
+        super().__init__(message)  # type: ignore
+        self.state = self._state_cls(message=self, *args, **kwargs)
 
 
 class FAIL(PrefectStateSignal):


Test patch:
diff --git a/tests/engine/test_signals.py b/tests/engine/test_signals.py
--- a/tests/engine/test_signals.py
+++ b/tests/engine/test_signals.py
@@ -1,3 +1,4 @@
+import datetime
 import pytest
 
 from prefect.engine.signals import (
@@ -33,6 +34,26 @@ def test_signals_create_states():
     assert exc.value.state.result is None
 
 
+def test_signals_pass_arguments_to_states():
+    with pytest.raises(PrefectStateSignal) as exc:
+        raise SUCCESS("you did it!", result=100)
+    assert exc.value.state.message is exc.value
+    assert exc.value.state.result == 100
+    assert str(exc.value.state.message) == "you did it!"
+
+
+def test_signals_dont_pass_invalid_arguments_to_states():
+    with pytest.raises(TypeError):
+        raise SUCCESS(bad_result=100)
+
+
+def test_retry_signals_can_set_retry_time():
+    date = datetime.datetime(2019, 1, 1)
+    with pytest.raises(PrefectStateSignal) as exc:
+        raise RETRY(start_time=date)
+    assert exc.value.state.start_time == date
+
+
 @pytest.mark.parametrize(
     "signal,state",
     [
@@ -50,3 +71,11 @@ def test_signals_creates_correct_states(signal, state):
     assert type(exc.value.state) is state
     assert exc.value.state.message is exc.value
     assert str(exc.value.state.message) == state.__name__
+
+
+def test_retry_signals_carry_default_retry_time_on_state():
+    with pytest.raises(Exception) as exc:
+        raise RETRY()
+    assert exc.value.state.start_time is not None
+    now = datetime.datetime.utcnow()
+    assert now - exc.value.state.start_time < datetime.timedelta(seconds=0.1)
diff --git a/tests/engine/test_task_runner.py b/tests/engine/test_task_runner.py
--- a/tests/engine/test_task_runner.py
+++ b/tests/engine/test_task_runner.py
@@ -134,6 +134,26 @@ def test_task_that_fails_gets_retried_up_to_1_time():
     assert isinstance(state, Failed)
 
 
+def test_task_that_raises_retry_has_start_time_recognized():
+    now = datetime.datetime.utcnow()
+
+    class RetryNow(Task):
+        def run(self):
+            raise signals.RETRY()
+
+    class Retry5Min(Task):
+        def run(self):
+            raise signals.RETRY(start_time=now + datetime.timedelta(minutes=5))
+
+    state = TaskRunner(task=RetryNow()).run()
+    assert isinstance(state, Retrying)
+    assert now - state.start_time < datetime.timedelta(seconds=0.1)
+
+    state = TaskRunner(task=Retry5Min()).run()
+    assert isinstance(state, Retrying)
+    assert state.start_time == now + datetime.timedelta(minutes=5)
+
+
 def test_task_that_raises_retry_gets_retried_even_if_max_retries_is_set():
     """
     Test that tasks that raise a retry signal get retried even if they exceed max_retries