Repository: PrefectHQ/prefect

Issue:
flow.visualize(flow_state=flow_state) doesn't work with mapped tasks
Simple example:
```python
from prefect import task, Flow

@task
def extract():
    return [0, 1, 2]

@task
def transform(x):
    return 1/x

with Flow() as f:
    e = extract()
    t = transform.map(e)

res = f.run(return_tasks=f.tasks)
f.visualize() # works just fine
f.visualize(flow_state=res) # raises ValueError
```

how we should color / represent the mapped task in this setting is up for debate, or whether we should simply raise a warning / informative error that colors are not supported with mapped tasks.


Discussion:

Resolving patch:
diff --git a/src/prefect/core/flow.py b/src/prefect/core/flow.py
--- a/src/prefect/core/flow.py
+++ b/src/prefect/core/flow.py
@@ -929,10 +929,14 @@ def visualize(self, flow_state: "prefect.engine.state.State" = None) -> object:
             )
             raise ImportError(msg)
 
-        def get_color(task: Task) -> str:
-            assert flow_state  # mypy assert
-            assert isinstance(flow_state.result, dict)  # mypy assert
-            state = flow_state.result.get(task)
+        def get_color(task: Task, map_index: int = None) -> str:
+            assert flow_state
+            assert isinstance(flow_state.result, dict)
+
+            if map_index is not None:
+                state = flow_state.result.get(task, [])[map_index]
+            else:
+                state = flow_state.result.get(task)
             if state is not None:
                 assert state is not None  # mypy assert
                 return state.color + "80"
@@ -944,18 +948,45 @@ def get_color(task: Task) -> str:
             is_mapped = any(edge.mapped for edge in self.edges_to(t))
             shape = "box" if is_mapped else "ellipse"
             name = "{} <map>".format(t.name) if is_mapped else t.name
-            kwargs = (
-                {}
-                if not flow_state
-                else dict(color=get_color(t), style="filled", colorscheme="svg")
-            )
-            graph.node(str(id(t)), name, shape=shape, **kwargs)
+            if is_mapped and flow_state:
+                assert isinstance(flow_state.result, dict)
+                for map_index, _ in enumerate(flow_state.result[t]):
+                    kwargs = dict(
+                        color=get_color(t, map_index=map_index),
+                        style="filled",
+                        colorscheme="svg",
+                    )
+                    graph.node(str(id(t)) + str(map_index), name, shape=shape, **kwargs)
+            else:
+                kwargs = (
+                    {}
+                    if not flow_state
+                    else dict(color=get_color(t), style="filled", colorscheme="svg")
+                )
+                graph.node(str(id(t)), name, shape=shape, **kwargs)
 
         for e in self.edges:
             style = "dashed" if e.mapped else None
-            graph.edge(
-                str(id(e.upstream_task)), str(id(e.downstream_task)), e.key, style=style
-            )
+            if (
+                e.mapped
+                or any(edge.mapped for edge in self.edges_to(e.downstream_task))
+            ) and flow_state:
+                assert isinstance(flow_state.result, dict)
+                assert isinstance(flow_state.result[e.downstream_task], list)
+                for map_index, _ in enumerate(flow_state.result[e.downstream_task]):
+                    graph.edge(
+                        str(id(e.upstream_task)),
+                        str(id(e.downstream_task)) + str(map_index),
+                        e.key,
+                        style=style,
+                    )
+            else:
+                graph.edge(
+                    str(id(e.upstream_task)),
+                    str(id(e.downstream_task)),
+                    e.key,
+                    style=style,
+                )
 
         try:
             from IPython import get_ipython
diff --git a/src/prefect/engine/task_runner.py b/src/prefect/engine/task_runner.py
--- a/src/prefect/engine/task_runner.py
+++ b/src/prefect/engine/task_runner.py
@@ -152,6 +152,7 @@ def run(
         executor = executor or DEFAULT_EXECUTOR
 
         # Initialize CloudHandler and get task run version
+        db_state = None
         if config.get("prefect_cloud", None):
             flow_run_id = context.get("flow_run_id", None)
             task_run_info = self.client.get_task_run_info(
@@ -160,10 +161,12 @@ def run(
                 map_index=map_index,
                 result_handler=self.result_handler,
             )
+            db_state = task_run_info.state  # type: ignore
             context.update(
                 _task_run_version=task_run_info.version,  # type: ignore
                 _task_run_id=task_run_info.id,  # type: ignore
             )
+        state = state or db_state or Pending()
 
         # construct task inputs
         task_inputs = {}  # type: Dict[str, Any]
@@ -202,12 +205,6 @@ def run(
         with prefect.context(context, _task_name=self.task.name, _map_index=map_index):
 
             try:
-                # determine starting state
-                new_state = None
-                for handler in self.state_handlers:
-                    new_state = handler(self, None, new_state)
-                state = new_state or state or Pending()
-
                 # retrieve the run number and place in context
                 state = self.get_run_count(state=state)
 
@@ -280,6 +277,14 @@ def run(
                 state = exc.state
                 state.cached_inputs = task_inputs or {}  # type: ignore
 
+            except Exception as exc:
+                state = Failed(
+                    message="Unexpected error while running Task", result=exc
+                )
+                raise_on_exception = prefect.context.get("_raise_on_exception", False)
+                if raise_on_exception:
+                    raise exc
+
             finally:  # resource is now available
                 for ticket, q in zip(tickets, queues):
                     q.put(ticket)


Test patch:
diff --git a/tests/client/test_client.py b/tests/client/test_client.py
--- a/tests/client/test_client.py
+++ b/tests/client/test_client.py
@@ -109,27 +109,15 @@ def test_client_token_priotizes_config_over_file(monkeypatch):
     assert client.token == "config-token"
 
 
-def test_client_logs_in_and_saves_token(monkeypatch):
-    post = MagicMock(
-        return_value=MagicMock(
-            ok=True, json=MagicMock(return_value=dict(token="secrettoken"))
-        )
-    )
+def test_client_doesnt_write_to_file_if_token_provided_from_config(monkeypatch):
     monkeypatch.setattr("os.path.exists", MagicMock(return_value=True))
     mock_file = mock_open()
     monkeypatch.setattr("builtins.open", mock_file)
-    monkeypatch.setattr("requests.post", post)
-    with set_temporary_config({"cloud.api": "http://my-cloud.foo"}):
+    with set_temporary_config(
+        {"cloud.api": "http://my-cloud.foo", "cloud.auth_token": "token"}
+    ):
         client = Client()
-    client.login("test@example.com", "1234")
-    assert post.called
-    assert post.call_args[0][0] == "http://my-cloud.foo/login_email"
-    assert post.call_args[1]["auth"] == ("test@example.com", "1234")
-    assert client.token == "secrettoken"
-    assert mock_file.call_args[0] == (
-        os.path.expanduser("~/.prefect/.credentials/auth_token"),
-        "w+",
-    )
+    assert not mock_file.called
 
 
 @pytest.mark.parametrize("cloud", [True, False])
diff --git a/tests/client/test_cloud.py b/tests/client/test_cloud.py
new file mode 100644
--- /dev/null
+++ b/tests/client/test_cloud.py
@@ -0,0 +1,165 @@
+import pytest
+from unittest.mock import MagicMock
+
+import prefect
+from prefect.client import Client
+from prefect.client.result_handlers import ResultHandler
+from prefect.engine.flow_runner import FlowRunner
+from prefect.engine.task_runner import TaskRunner
+from prefect.engine.state import (
+    Failed,
+    Running,
+    Pending,
+    Success,
+    Finished,
+    TriggerFailed,
+    TimedOut,
+    Skipped,
+)
+from prefect.utilities.configuration import set_temporary_config
+
+
+@pytest.fixture(autouse=True)
+def cloud_settings():
+    with set_temporary_config(
+        {
+            "cloud.api": "http://my-cloud.foo",
+            "prefect_cloud": True,
+            "cloud.auth_token": "token",
+        }
+    ):
+        yield
+
+
+class TestFlowRunner:
+    def test_flow_runner_calls_client_the_approriate_number_of_times(self, monkeypatch):
+        flow = prefect.Flow(name="test")
+        get_flow_run_info = MagicMock(return_value=MagicMock(state=None))
+        set_flow_run_state = MagicMock()
+        client = MagicMock(
+            get_flow_run_info=get_flow_run_info, set_flow_run_state=set_flow_run_state
+        )
+        monkeypatch.setattr(
+            "prefect.engine.flow_runner.Client", MagicMock(return_value=client)
+        )
+        res = FlowRunner(flow=flow).run()
+
+        ## assertions
+        assert get_flow_run_info.call_count == 1  # one time to pull latest state
+        assert set_flow_run_state.call_count == 2  # Pending -> Running -> Success
+
+        states = [call[1]["state"] for call in set_flow_run_state.call_args_list]
+        assert states == [Running(), Success(result=dict())]
+
+    @pytest.mark.parametrize(
+        "state", [Finished, Success, Skipped, Failed, TimedOut, TriggerFailed]
+    )
+    def test_flow_runner_respects_the_db_state(self, monkeypatch, state):
+        flow = prefect.Flow(name="test")
+        db_state = state("already", result=10)
+        get_flow_run_info = MagicMock(return_value=MagicMock(state=db_state))
+        set_flow_run_state = MagicMock()
+        client = MagicMock(
+            get_flow_run_info=get_flow_run_info, set_flow_run_state=set_flow_run_state
+        )
+        monkeypatch.setattr(
+            "prefect.engine.flow_runner.Client", MagicMock(return_value=client)
+        )
+        res = FlowRunner(flow=flow).run()
+
+        ## assertions
+        assert get_flow_run_info.call_count == 1  # one time to pull latest state
+        assert set_flow_run_state.call_count == 0  # never needs to update state
+        assert res == db_state
+
+    @pytest.mark.parametrize(
+        "state", [Finished, Success, Skipped, Failed, TimedOut, TriggerFailed]
+    )
+    def test_flow_runner_prioritizes_kwarg_states_over_db_states(
+        self, monkeypatch, state
+    ):
+        flow = prefect.Flow(name="test")
+        db_state = state("already", result=10)
+        get_flow_run_info = MagicMock(return_value=MagicMock(state=db_state))
+        set_flow_run_state = MagicMock()
+        client = MagicMock(
+            get_flow_run_info=get_flow_run_info, set_flow_run_state=set_flow_run_state
+        )
+        monkeypatch.setattr(
+            "prefect.engine.flow_runner.Client", MagicMock(return_value=client)
+        )
+        res = FlowRunner(flow=flow).run(state=Pending("let's do this"))
+
+        ## assertions
+        assert get_flow_run_info.call_count == 1  # one time to pull latest state
+        assert set_flow_run_state.call_count == 2  # Pending -> Running -> Success
+
+        states = [call[1]["state"] for call in set_flow_run_state.call_args_list]
+        assert states == [Running(), Success(result=dict())]
+
+
+class TestTaskRunner:
+    def test_task_runner_calls_client_the_approriate_number_of_times(self, monkeypatch):
+        task = prefect.Task(name="test")
+        get_task_run_info = MagicMock(return_value=MagicMock(state=None))
+        set_task_run_state = MagicMock()
+        client = MagicMock(
+            get_task_run_info=get_task_run_info, set_task_run_state=set_task_run_state
+        )
+        monkeypatch.setattr(
+            "prefect.engine.task_runner.Client", MagicMock(return_value=client)
+        )
+        res = TaskRunner(task=task).run()
+
+        ## assertions
+        assert get_task_run_info.call_count == 1  # one time to pull latest state
+        assert set_task_run_state.call_count == 2  # Pending -> Running -> Success
+
+        states = [call[1]["state"] for call in set_task_run_state.call_args_list]
+        assert states == [Running(), Success()]
+
+    @pytest.mark.parametrize(
+        "state", [Finished, Success, Skipped, Failed, TimedOut, TriggerFailed]
+    )
+    def test_task_runner_respects_the_db_state(self, monkeypatch, state):
+        task = prefect.Task(name="test")
+        db_state = state("already", result=10)
+        get_task_run_info = MagicMock(return_value=MagicMock(state=db_state))
+        set_task_run_state = MagicMock()
+        client = MagicMock(
+            get_task_run_info=get_task_run_info, set_task_run_state=set_task_run_state
+        )
+        monkeypatch.setattr(
+            "prefect.engine.task_runner.Client", MagicMock(return_value=client)
+        )
+        res = TaskRunner(task=task).run()
+
+        ## assertions
+        assert get_task_run_info.call_count == 1  # one time to pull latest state
+        assert set_task_run_state.call_count == 0  # never needs to update state
+        assert res == db_state
+
+    @pytest.mark.parametrize(
+        "state", [Finished, Success, Skipped, Failed, TimedOut, TriggerFailed]
+    )
+    def test_task_runner_prioritizes_kwarg_states_over_db_states(
+        self, monkeypatch, state
+    ):
+        task = prefect.Task(name="test")
+        db_state = state("already", result=10)
+        get_task_run_info = MagicMock(return_value=MagicMock(state=db_state))
+        set_task_run_state = MagicMock()
+        client = MagicMock(
+            get_task_run_info=get_task_run_info, set_task_run_state=set_task_run_state
+        )
+        monkeypatch.setattr(
+            "prefect.engine.task_runner.Client", MagicMock(return_value=client)
+        )
+        res = TaskRunner(task=task).run(state=Pending("let's do this"))
+
+        ## assertions
+        assert get_task_run_info.call_count == 1  # one time to pull latest state
+        assert set_task_run_state.call_count == 2  # Pending -> Running -> Success
+
+        states = [call[1]["state"] for call in set_task_run_state.call_args_list]
+        assert states == [Running(), Success()]
diff --git a/tests/core/test_flow.py b/tests/core/test_flow.py
--- a/tests/core/test_flow.py
+++ b/tests/core/test_flow.py
@@ -12,6 +12,7 @@
 from prefect.core.flow import Flow
 from prefect.core.task import Parameter, Task
 from prefect.engine.signals import PrefectError
+from prefect.engine.state import Success, Failed, Skipped
 from prefect.tasks.core.function import FunctionTask
 from prefect.utilities.tasks import task, unmapped
 from prefect.utilities.configuration import set_temporary_config
@@ -845,6 +846,49 @@ def test_viz_reflects_mapping(self):
         assert "label=x style=dashed" in graph.source
         assert "label=y style=dashed" in graph.source
 
+    @pytest.mark.parametrize("state", [Success(), Failed(), Skipped()])
+    def test_viz_if_flow_state_provided(self, state):
+        import graphviz
+
+        ipython = MagicMock(
+            get_ipython=lambda: MagicMock(config=dict(IPKernelApp=True))
+        )
+        with patch.dict("sys.modules", IPython=ipython):
+            t = Task(name="a_nice_task")
+            f = Flow()
+            f.add_task(t)
+            graph = f.visualize(flow_state=Success(result={t: state}))
+        assert "label=a_nice_task" in graph.source
+        assert 'color="' + state.color + '80"' in graph.source
+        assert "shape=ellipse" in graph.source
+
+    def test_viz_reflects_mapping_if_flow_state_provided(self):
+        ipython = MagicMock(
+            get_ipython=lambda: MagicMock(config=dict(IPKernelApp=True))
+        )
+        add = AddTask(name="a_nice_task")
+        list_task = Task(name="a_list_task")
+
+        with patch.dict("sys.modules", IPython=ipython):
+            with Flow() as f:
+                res = add.map(x=list_task, y=8)
+            graph = f.visualize(
+                flow_state=Success(
+                    result={res: [Success(), Failed()], list_task: Success()}
+                )
+            )
+
+        # one colored node for each mapped result
+        assert 'label="a_nice_task <map>" color="#00800080"' in graph.source
+        assert 'label="a_nice_task <map>" color="#FF000080"' in graph.source
+        assert 'label=a_list_task color="#00800080"' in graph.source
+        assert 'label=8 color="#00000080"' in graph.source
+
+        # two edges for each input to add()
+        for var in ["x", "y"]:
+            for index in [0, 1]:
+                assert "{0} [label={1} style=dashed]".format(index, var) in graph.source
+
     @pytest.mark.parametrize(
         "error",
         [
diff --git a/tests/engine/test_flow_runner.py b/tests/engine/test_flow_runner.py
--- a/tests/engine/test_flow_runner.py
+++ b/tests/engine/test_flow_runner.py
@@ -1024,3 +1024,23 @@ def handler(flow, old, new):
         flow = Flow(state_handlers=[handler])
         state = FlowRunner(flow=flow).run()
         assert state.is_failed()
+
+
+def test_improper_use_of_unmapped_fails_gracefully():
+    add = AddTask()
+    x = Parameter("x", default=[1, 2, 3])
+    with Flow() as f:
+        res = add.map(x, y=8)  # incorrect, should use `unmapped`
+
+    state = FlowRunner(flow=f).run(return_tasks=f.tasks)
+    assert state.is_failed()
+
+    # make sure tasks were still returned with the correct states
+    x_state = state.result.pop(x)
+    res_state = state.result.pop(res)
+    y_state = state.result.popitem()[1]
+    assert x_state.is_successful()
+    assert x_state.result == [1, 2, 3]
+    assert y_state.is_successful()
+    assert y_state.result == 8
+    assert res_state.is_failed()
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
@@ -989,8 +989,8 @@ def test_task_handler_that_doesnt_return_state(self):
 class TestTaskRunnerStateHandlers:
     def test_task_runner_handlers_are_called(self):
         TaskRunner(task=Task(), state_handlers=[task_runner_handler]).run()
-        # the task changed state three times: Initialization -> Pending -> Running -> Success
-        assert handler_results["TaskRunner"] == 3
+        # the task changed state three times: Pending -> Running -> Success
+        assert handler_results["TaskRunner"] == 2
 
     def test_task_runner_handlers_are_called_on_retry(self):
         @prefect.task(max_retries=1, retry_delay=timedelta(0))
@@ -998,15 +998,15 @@ def fn():
             1 / 0
 
         TaskRunner(task=fn, state_handlers=[task_runner_handler]).run()
-        # the task changed state four times: Initialization -> Pending -> Running -> Failed -> Retry
-        assert handler_results["TaskRunner"] == 4
+        # the task changed state four times: Pending -> Running -> Failed -> Retry
+        assert handler_results["TaskRunner"] == 3
 
     def test_multiple_task_runner_handlers_are_called(self):
         TaskRunner(
             task=Task(), state_handlers=[task_runner_handler, task_runner_handler]
         ).run()
-        # each task changed state three times: Initialization -> Pending -> Running -> Success
-        assert handler_results["TaskRunner"] == 6
+        # each task changed state three times: Pending -> Running -> Success
+        assert handler_results["TaskRunner"] == 4
 
     def test_multiple_task_runner_handlers_are_called_in_sequence(self):
         # the second task handler will assert the result of the first task handler is a state
@@ -1203,3 +1203,17 @@ def run(self):
     runner = TaskRunner(t2)
     out = runner.run(upstream_states={e: Success(result=1)}, context=dict(resume=True))
     assert out.is_successful()
+
+
+@pytest.mark.parametrize("mapped", [False, True])
+def test_improperly_mapped_edge_fails_gracefully(mapped):
+    add = AddTask()
+    x = Task()
+    e = Edge(x, add, mapped=True, key="x")
+    f = Edge(x, 8, mapped=True, key="y")
+
+    state = TaskRunner(add).run(
+        upstream_states={e: Success(result=[1, 2, 3]), f: Success(result=8)},
+        mapped=mapped,
+    )
+    assert state.is_failed()
diff --git a/tests/utilities/test_tests.py b/tests/utilities/test_tests.py
--- a/tests/utilities/test_tests.py
+++ b/tests/utilities/test_tests.py
@@ -82,11 +82,16 @@ def get_task_run_state(self, *args, **kwargs):
             raise RuntimeError("I represent bad code in the task runner.")
 
     flow = Flow()
-    flow.add_task(MathTask())
-    flow_state = FlowRunner(flow=flow, task_runner_cls=BadTaskRunner).run()
+    math_task = MathTask()
+    flow.add_task(math_task)
+    flow_state = FlowRunner(flow=flow, task_runner_cls=BadTaskRunner).run(
+        return_tasks=[math_task]
+    )
     assert isinstance(flow_state, state.Failed)
-    assert isinstance(flow_state.result, RuntimeError)
-    assert "I represent bad code in the task runner." == str(flow_state.result)
+    assert flow_state.result[math_task].is_failed()
+    caught_exc = flow_state.result[math_task].result
+    assert isinstance(caught_exc, RuntimeError)
+    assert "I represent bad code in the task runner." == str(caught_exc)
 
 
 def test_raise_on_exception_plays_well_with_context():