class TestMapTask(object): def setup(self): self.args = range(3) self.pool = Mock() self.task = MapTask(add, self.args, self.pool) def test_iter_input(self): def args(data): for i in range(3): yield (data, i) self.task = MapTask(add, args, self.pool) iter_input = self.task.iter_input('foo') assert list(iter_input) == [ ('foo', 0), ('foo', 1), ('foo', 2), ] def test_calls_map_on_pool(self): self.task.iter_input = Mock() self.task('foo') self.pool.map.assert_called_with(add, self.task.iter_input()) def test_map_task_uses_thread_pool_default(self): tsk = MapTask(self.args, add) assert isinstance(tsk.pool, ThreadPool)
def test_iter_input(self): def args(data): for i in range(3): yield (data, i) self.task = MapTask(add, args, self.pool) iter_input = self.task.iter_input('foo') assert list(iter_input) == [ ('foo', 0), ('foo', 1), ('foo', 2), ]
def dig_it(zone): cmd = 'dig +short %s' % zone print('running: %s' % cmd.strip()) ips = subprocess.check_output(cmd.split()).strip().split('\n') return (zone, ips) myflow = [ foo, IfTask(check, [bar], [baz]), MapTask(dig_it, [ 'ionrock.org', 'google.com', 'yahoo.com', ]), finish, ] def main(): results = do_flow(myflow, {}) pprint(results) print('Done') if __name__ == '__main__': main()
def setup(self): self.args = range(3) self.pool = Mock() self.task = MapTask(add, self.args, self.pool)
def test_map_task_uses_thread_pool_default(self): tsk = MapTask(self.args, add) assert isinstance(tsk.pool, ThreadPool)