def test_upstream_and_downstream_error_msgs_when_task_is_not_in_flow():
f = Flow(name="test")
t = Task()
with pytest.raises(ValueError) as e:
f.edges_to(t)
assert "was not found in Flow" in e
with pytest.raises(ValueError) as e:
f.edges_from(t)
assert "was not found in Flow" in e
with pytest.raises(ValueError) as e:
f.upstream_tasks(t)
assert "was not found in Flow" in e
with pytest.raises(ValueError) as e:
f.downstream_tasks(t)
assert "was not found in Flow" in e
def test_add_edge():
f = Flow(name="test")
t1 = Task()
t2 = Task()
f.add_edge(upstream_task=t1, downstream_task=t2)
assert f.upstream_tasks(t2) == set([t1])
assert f.upstream_tasks(t1) == set()
assert f.downstream_tasks(t2) == set()
assert f.downstream_tasks(t1) == set([t2])
assert f.edges_to(t2) == f.edges_from(t1)