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)