def test_deserialize_edges():
"""
Tests that edges are appropriately deserialized, even in they involve keys.
Also tests that tasks are deserialized in a way that reuses them in edges -- in other
words, when edges are loaded they use their corresponding task IDs to access the
correct Task objects out of a cache.
"""
class ArgTask(Task):
def run(self, x):
return x
f = Flow(name="test")
t1, t2, t3 = Task("a"), Task("b"), ArgTask("c")
f.add_edge(t1, t2)
f.add_edge(t2, t3, key="x")
f.add_edge(t1, t3, mapped=True)
serialized = f.serialize()
deserialized = FlowSchema().load(serialized)
d1, d2, d3 = sorted(deserialized.tasks, key=lambda t: t.name)
assert deserialized.edges == {
Edge(d1, d2),
Edge(d2, d3, key="x"),
Edge(d1, d3, mapped=True),
}
def test_reference_tasks():
x = Task("x")
y = Task("y")
z = Task("z")
f = Flow(name="test", tasks=[x, y, z])
f.set_reference_tasks([y])
assert f.reference_tasks() == {y}
f2 = FlowSchema().load(f.serialize())
assert f2.reference_tasks() == {t for t in f2.tasks if t.name == "y"}
def test_serialize_flow_sorts_nested_schemas():
a = Parameter("a", default=1)
b = Parameter("b", default=2)
c = Task("c")
d = Task("d")
f = Flow("test")
d.set_upstream(c, flow=f)
c.set_upstream(b, flow=f).set_upstream(a, flow=f)
f.set_reference_tasks([d, c])
# Must use `f.serialize` instead of `FlowSchema().dump` because task slugs
# are not guaranteed to be set yet
serialized = f.serialize()
assert [param["slug"] for param in serialized["parameters"]] == ["a", "b"]
assert [task["slug"] for task in serialized["tasks"]] == ["a", "b", "c-1", "d-1"]
assert [
(edge["upstream_task"]["slug"], edge["downstream_task"]["slug"])
for edge in serialized["edges"]
] == [("a", "c-1"), ("b", "c-1"), ("c-1", "d-1")]
assert [task["slug"] for task in serialized["reference_tasks"]] == ["c-1", "d-1"]
def test_deserialize_serialized_flow_after_build(tmpdir):
flow = Flow(name="test",
storage=prefect.environments.storage.Local(tmpdir))
serialized_flow = flow.serialize(build=True)
deserialized = FlowSchema().load(serialized_flow)
assert isinstance(deserialized, Flow)
def test_deserialize_tasks():
tasks = [Task(n) for n in ["a", "b", "c"]]
f = Flow(name="test", tasks=tasks)
serialized = f.serialize()
deserialized = FlowSchema().load(serialized)
assert len(deserialized.tasks) == len(f.tasks)
def test_deserialize_serialized_flow_after_build():
flow = Flow(name="test",
environment=prefect.environments.LocalEnvironment())
serialized_flow = flow.serialize(build=True)
deserialized = FlowSchema().load(serialized_flow)
assert isinstance(deserialized, Flow)
