def test_branch(self):
class Inc(Runnable):
def next(self, state):
return state.updated(cnt=state.cnt + 1)
states = States(State(cnt=1), State(cnt=2))
branch = Map(Inc()) | ArgMin('cnt')
result = branch.run(states).result()
self.assertEqual(result.cnt, states[0].cnt + 1)
def test_map(self):
class A(Runnable, traits.ProblemDecomposer):
def next(self, state):
return state.updated(subproblem=state.problem)
a = A()
m = Map(a)
r = m.run(States(State(problem=1))).result()
self.assertSetEqual(m.inputs, a.inputs)
self.assertSetEqual(m.outputs, a.outputs)