Repository: PrefectHQ/prefect

Issue:
flow.visualize(flow_state) with double layer mapping
it appears that `flow.visualize()` with the flow_state kwarg and a flow with two sequential mapping layers provides a bad picture; I think the graph node ID isn't being looked up correctly somewhere
flow.visualize(flow_state) with double layer mapping
it appears that `flow.visualize()` with the flow_state kwarg and a flow with two sequential mapping layers provides a bad picture; I think the graph node ID isn't being looked up correctly somewhere


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
@@ -979,8 +979,11 @@ def get_color(task: Task, map_index: int = None) -> str:
                 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]):
+                    upstream_id = str(id(e.upstream_task))
+                    if any(edge.mapped for edge in self.edges_to(e.upstream_task)):
+                        upstream_id += str(map_index)
                     graph.edge(
-                        str(id(e.upstream_task)),
+                        upstream_id,
                         str(id(e.downstream_task)) + str(map_index),
                         e.key,
                         style=style,


Test patch:
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
@@ -14,7 +14,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.engine.state import Success, Failed, Skipped, TriggerFailed
 from prefect.tasks.core.function import FunctionTask
 from prefect.utilities.tasks import task, unmapped
 from prefect.utilities.configuration import set_temporary_config
@@ -853,6 +853,34 @@ def test_viz_reflects_mapping_if_flow_state_provided(self):
             for index in [0, 1]:
                 assert "{0} [label={1} style=dashed]".format(index, var) in graph.source
 
+    def test_viz_reflects_multiple_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:
+                first_res = add.map(x=list_task, y=8)
+                res = first_res.map(x=first_res, y=9)
+            graph = f.visualize(
+                flow_state=Success(
+                    result={
+                        res: [Success(), TriggerFailed()],
+                        list_task: Success(),
+                        first_res: [Success(), Failed()],
+                    }
+                )
+            )
+
+        assert "{first} -> {second} [label=x style=dashed]".format(
+            first=str(id(first_res)) + "0", second=str(id(res)) + "0"
+        )
+        assert "{first} -> {second} [label=x style=dashed]".format(
+            first=str(id(first_res)) + "1", second=str(id(res)) + "1"
+        )
+
     @pytest.mark.parametrize(
         "error",
         [