Add tests + fixes for op::flat_map