def test_visualize_raises_informative_importerror_without_graphviz(
self, monkeypatch):
f = Flow(name="test")
f.add_task(Task())
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
ipython = MagicMock(
get_ipython=lambda: MagicMock(config=dict(IPKernelApp=True)))
with patch.dict("sys.modules", IPython=ipython):
f = Flow(name="test")
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_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(name="test")
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_saves_graph_object_if_filename(self):
import graphviz
f = Flow(name="test")
f.add_task(Task(name="a_nice_task"))
with tempfile.NamedTemporaryFile() as tmp:
graph = f.visualize(filename=tmp.name)
with open(tmp.name, "r") as f:
contents = f.read()
assert "label=a_nice_task" in contents
assert "shape=ellipse" in contents
