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