Repository: PrefectHQ/prefect

Issue:
Add "raise_on_fail" context manager
Task runners and flow runners trap all errors, which can make debugging difficult. 

We could create a context manager that would raise errors rather than trapping them. Given a Flow `flow` that contains a task that raises some error:

```python
flow.run() # returns State(FAILED)

with raise_on_fail():
    flow.run() # actually raises an error at the appropriate place
```




Discussion:

Resolving patch:
diff --git a/prefect/engine/flow_runner.py b/prefect/engine/flow_runner.py
--- a/prefect/engine/flow_runner.py
+++ b/prefect/engine/flow_runner.py
@@ -50,6 +50,8 @@ def run(
 
             with self.executor.start():
 
+                raise_on_fail = prefect.context.get("_raise_on_fail", False)
+
                 try:
                     state = self._run(
                         state=state,
@@ -62,6 +64,8 @@ def run(
                     self.logger.info("Flow run DONTRUN")
                     # set state but no need to go through the executor
                     state = type(state)(return_task_states)
+                    if raise_on_fail:
+                        raise e
                 except Exception as e:
                     self.logger.info("Flow run FAIL")
                     # set state through executor
@@ -69,6 +73,8 @@ def run(
                     state = self.executor.set_state(
                         state, Failed, data=return_task_states
                     )
+                    if raise_on_fail:
+                        raise e
 
         return state
 
diff --git a/prefect/engine/task_runner.py b/prefect/engine/task_runner.py
--- a/prefect/engine/task_runner.py
+++ b/prefect/engine/task_runner.py
@@ -61,6 +61,8 @@ def run(
             # prepare executor
             with self.executor.start():
 
+                raise_on_fail = prefect.context.get("_raise_on_fail", False)
+
                 try:
                     state = self._run(
                         state=state,
@@ -77,7 +79,9 @@ def run(
                     state = self.executor.set_state(state, Success)
 
                 except signals.FAIL as e:
-                    state = self.handle_fail(state, data=dict(message=e))
+                    state = self.handle_fail(
+                        state, fail_mode=e, raise_on_fail=raise_on_fail
+                    )
 
                 except signals.RETRY:
                     state = self.handle_retry(state)
@@ -88,7 +92,9 @@ def run(
 
                 except Exception as e:
                     logging.info("Unexpected error while running task.")
-                    state = self.handle_fail(state, data=dict(message=e))
+                    state = self.handle_fail(
+                        state, fail_mode=e, raise_on_fail=raise_on_fail
+                    )
 
         return state
 
@@ -155,18 +161,28 @@ def _run(
 
         return state
 
-    def handle_fail(self, state, data=None):
+    def handle_fail(self, state, fail_mode=None, raise_on_fail=False):
         """
         Checks if a task is eligable for retry; otherwise marks it failed.
         """
         self.logger.info("Task FAILED")
-        if "Trigger failed" in str(data.get("message")):
-            return self.executor.set_state(state, TriggerFailed, data=data)
+        if "Trigger failed" in str(fail_mode):
+            state = self.executor.set_state(
+                state, TriggerFailed, data=dict(message=fail_mode)
+            )
+            if raise_on_fail:
+                raise fail_mode
+            else:
+                return state
         run_number = prefect.context.get("_task_run_number", 1)
         if run_number and run_number <= self.task.max_retries:
             return self.handle_retry(state)
         else:
-            return self.executor.set_state(state, Failed, data=data)
+            state = self.executor.set_state(state, Failed, data=dict(message=fail_mode))
+            if raise_on_fail:
+                raise fail_mode
+            else:
+                return state
 
     def handle_retry(self, state, retry_time=None):
         # TODO exponential backoff based on run_number


Test patch:
diff --git a/prefect/utilities/tests.py b/prefect/utilities/tests.py
--- a/prefect/utilities/tests.py
+++ b/prefect/utilities/tests.py
@@ -26,6 +26,12 @@ def set_config(keys, value):
         prefect.config.__dict__.update(old_config)
 
 
+@contextmanager
+def raise_on_fail():
+    with prefect.context(_raise_on_fail=True):
+        yield
+
+
 def run_flow_runner_test(
     flow: Flow,
     expected_state: str,
diff --git a/tests/utilities/test_tests.py b/tests/utilities/test_tests.py
new file mode 100644
--- /dev/null
+++ b/tests/utilities/test_tests.py
@@ -0,0 +1,77 @@
+import pytest
+
+import prefect
+
+from prefect.utilities.tests import raise_on_fail
+from prefect.core import Flow, Task
+from prefect.engine import TaskRunner
+
+
+class BusinessTask(Task):
+    def run(self):
+        raise prefect.signals.FAIL("needs more blockchain!")
+
+
+class MathTask(Task):
+    def run(self):
+        1 / 0
+
+
+def test_raise_on_fail_raises_basic_error():
+    flow = Flow()
+    flow.add_task(MathTask())
+    with pytest.raises(ZeroDivisionError):
+        with raise_on_fail():
+            flow.run()
+
+
+def test_raise_on_fail_raises_basic_prefect_signal():
+    flow = Flow()
+    flow.add_task(BusinessTask())
+    with pytest.raises(prefect.signals.FAIL) as error:
+        with raise_on_fail():
+            flow.run()
+    assert "needs more blockchain!" in str(error)
+
+
+def test_raise_on_fail_works_at_the_task_level_with_error():
+    taskrunner = TaskRunner(task=MathTask())
+    with pytest.raises(ZeroDivisionError):
+        with raise_on_fail():
+            taskrunner.run()
+
+
+def test_raise_on_fail_works_at_the_task_level_with_signal():
+    taskrunner = TaskRunner(task=BusinessTask())
+    with pytest.raises(prefect.signals.FAIL) as error:
+        with raise_on_fail():
+            taskrunner.run()
+    assert "needs more blockchain!" in str(error)
+
+
+def test_core_code_errors_bubble_up(monkeypatch):
+    flow = Flow()
+    flow.add_task(MathTask())
+
+    class BadTaskRunner(TaskRunner):
+        def handle_fail(self, *args, **kwargs):
+            raise RuntimeError("I'm not cool with this.")
+
+    monkeypatch.setattr(prefect.engine, "TaskRunner", BadTaskRunner)
+    with pytest.raises(RuntimeError) as error:
+        with raise_on_fail():
+            flow.run()
+    assert "I'm not cool with this." in str(error)
+
+
+def test_raise_on_fail_raises_basic_error():
+    flow = Flow()
+    flow.add_task(MathTask())
+    try:
+        assert "_raise_on_fail" not in prefect.context
+        with raise_on_fail():
+            assert "_raise_on_fail" in prefect.context
+            flow.run()
+        assert "_raise_on_fail" not in prefect.context
+    except ZeroDivisionError:
+        pass