def test_setting_reference_tasks_clears_cache(self):
f = Flow(name="test")
t1 = Task()
f.add_task(t1)
f._cache[1] = 2
f.set_reference_tasks([t1])
assert 1 not in f._cache
def test_flow_parameters():
f = Flow(name="test")
x = Parameter("x")
y = Parameter("y", default=1)
f.add_task(x)
f.add_task(y)
assert f.parameters() == {x, y}
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_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_validate_missing_reference_tasks():
f = Flow(name="test")
t1 = Task()
t2 = Task()
f.add_task(t1)
f.add_task(t2)
f.set_reference_tasks([t1])
f.tasks.remove(t1)
with pytest.raises(ValueError) as exc:
f.validate()
assert "reference tasks are not contained" in str(exc.value).lower()
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
def test_equality_based_on_tasks(self):
f1 = Flow(name="test")
f2 = Flow(name="test")
t1 = Task()
t2 = Task()
t3 = Task()
for f in [f1, f2]:
f.add_task(t1)
f.add_task(t2)
assert f1 == f2
f2.add_task(t3)
assert f1 != f2
def test_adding_a_task_to_a_flow_twice_is_ok():
f = Flow(name="test")
t = Task()
f.add_task(t)
f.add_task(t)
def test_add_task_raise_an_error_if_the_task_is_not_a_task_class():
f = Flow(name="test")
with pytest.raises(TypeError):
f.add_task(1)
def test_add_task_returns_task():
f = Flow(name="test")
t = Task()
t2 = f.add_task(t)
assert t2 is t
def test_add_task_to_flow():
f = Flow(name="test")
t = Task()
f.add_task(t)
assert t in f.tasks
def test_adding_task_clears_cache(self):
f = Flow(name="test")
f._cache[1] = 2
f.add_task(Task())
assert 1 not in f._cache
def test_raise_error_if_two_parameters_have_same_name():
f = Flow(name="test")
f.add_task(Parameter("x"))
assert "x" in {p.name for p in f.parameters()}
with pytest.raises(ValueError):
f.add_task(Parameter("x"))