def test_apply_map_simple(self, api):
if api == "functional":
def func(x, y, z):
a = add(x, y)
a.name = "add-a"
b = add(a, z)
b.name = "add-b"
c = add(b, 1)
c.name = "add-c"
return inc(c)
with Flow("test") as flow:
y = ranged(3)
z = edges.unmapped(1)
res = apply_map(func, range(3, 6), y=y, z=z)
else:
def func(x, y, z, flow):
a = add.copy(name="add-a").bind(x, y, flow=flow)
b = add.copy(name="add-b").bind(a, z, flow=flow)
c = add.copy(name="add-c").bind(b, 1, flow=flow)
return inc.copy().bind(c, flow=flow)
flow = Flow("test")
y = ranged.copy().bind(3, flow=flow)
z = edges.unmapped(tasks.as_task(1, flow=flow))
res = apply_map(func, range(3, 6), y=y, z=z, flow=flow)
consts = {t.name: c for t, c in flow.constants.items()}
assert consts == {
"ranged": {
"n": 3
},
"add-b": {
"y": 1
},
"add-c": {
"y": 1
}
}
for task in flow.tasks:
if task.name != "ranged":
for e in flow.edges_to(task):
assert e.mapped
state = flow.run()
assert state.result[res].result == [6, 8, 10]
