Repository: PrefectHQ/prefect

Issue:
Add a few `flow.visualize()` unit tests
We should have a few tests just to make sure the colors work, etc.

Adding this as a small, low priority issue for picking up whenever any of us need some added diversity of tasks


Discussion:
Note: should also change the `except NameError` within `flow.visualize()` to `except ImportError`

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
@@ -986,7 +986,7 @@ def get_color(task: Task) -> str:
 
             if get_ipython().config.get("IPKernelApp") is not None:
                 return graph
-        except NameError:
+        except Exception:
             pass
 
         with tempfile.NamedTemporaryFile() as tmp:


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
@@ -3,6 +3,8 @@
 
 import cloudpickle
 import pytest
+import sys
+from unittest.mock import MagicMock, patch
 
 import prefect
 from prefect.core.edge import Edge
@@ -869,18 +871,63 @@ def test_build_environment():
     assert isinstance(key, prefect.environments.LocalEnvironment)
 
 
-def test_visualize_raises_informative_importerror_without_graphviz(monkeypatch):
-    f = Flow()
-    f.add_task(Task())
+class TestFlowVisualize:
+    def test_visualize_raises_informative_importerror_without_graphviz(
+        self, monkeypatch
+    ):
+        f = Flow()
+        f.add_task(Task())
 
-    import sys
+        with monkeypatch.context() as m:
+            m.setattr(sys, "path", "")
+            with pytest.raises(ImportError) as exc:
+                f.visualize()
 
-    with monkeypatch.context() as m:
-        m.setattr(sys, "path", "")
-        with pytest.raises(ImportError) as exc:
-            f.visualize()
+        assert "pip install prefect[viz]" in repr(exc.value)
+
+    def test_viz_returns_graph_object_if_in_ipython(self):
+        import graphviz
 
-    assert "pip install prefect[viz]" in repr(exc.value)
+        ipython = MagicMock(
+            get_ipython=lambda: MagicMock(config=dict(IPKernelApp=True))
+        )
+        with patch.dict("sys.modules", IPython=ipython):
+            f = Flow()
+            f.add_task(Task(name="a_nice_task"))
+            graph = f.visualize()
+        assert "label=a_nice_task" in graph.source
+        assert "shape=ellipse" in graph.source
+
+    def test_viz_reflects_mapping(self):
+        ipython = MagicMock(
+            get_ipython=lambda: MagicMock(config=dict(IPKernelApp=True))
+        )
+        with patch.dict("sys.modules", IPython=ipython):
+            with Flow() as f:
+                res = AddTask(name="a_nice_task").map(x=Task(name="a_list_task"), y=8)
+            graph = f.visualize()
+        assert 'label="a_nice_task <map>" shape=box' in graph.source
+        assert "label=a_list_task shape=ellipse" in graph.source
+        assert "label=x style=dashed" in graph.source
+        assert "label=y style=dashed" in graph.source
+
+    @pytest.mark.parametrize(
+        "error",
+        [
+            ImportError("abc"),
+            ValueError("abc"),
+            TypeError("abc"),
+            NameError("abc"),
+            AttributeError("abc"),
+        ],
+    )
+    def test_viz_renders_if_ipython_isnt_installed_or_errors(self, error):
+        graphviz = MagicMock()
+        ipython = MagicMock(get_ipython=MagicMock(side_effect=error))
+        with patch.dict("sys.modules", graphviz=graphviz, IPython=ipython):
+            with Flow() as f:
+                res = AddTask(name="a_nice_task").map(x=Task(name="a_list_task"), y=8)
+            f.visualize()
 
 
 class TestCache: