test/array_test.rb in rumonade-0.4.0 vs test/array_test.rb in rumonade-0.4.1

@@ -22,10 +22,11 @@
     end
   end
 
   def test_flat_map_behaves_correctly
     assert_equal ["FOO", "BAR"], ["foo", "bar"].flat_map { |s| [s.upcase] }
+    assert_equal [2, 4, 6], [1, 2, 3].flat_map { |i| i * 2 }
   end
 
   def test_map_behaves_correctly
     assert_equal ["FOO", "BAR"], ["foo", "bar"].map { |s| s.upcase }
   end
@@ -33,13 +34,25 @@
   def test_shallow_flatten_behaves_correctly
     assert_equal [0, 1, [2], [[3]], [[[4]]]], [0, [1], [[2]], [[[3]]], [[[[4]]]]].shallow_flatten
     assert_equal [1], [None, Some(1)].shallow_flatten
     assert_equal [1, Some(2)], [None, Some(1), Some(Some(2))].shallow_flatten
     assert_equal [Some(Some(None))], [Some(Some(Some(None)))].shallow_flatten
+    assert_equal [Some(1), Some(None), None], [Some(Some(1)), Some(Some(None)), [None]].shallow_flatten
   end
 
   def test_flatten_behaves_correctly
     assert_equal [0, 1, 2, 3, 4], [0, [1], [[2]], [[[3]]], [[[[4]]]]].flatten
     assert_equal [1, 2], [None, Some(1), Some(Some(2))].flatten
     assert_equal [], [Some(Some(Some(None)))].flatten
+  end
+
+  def test_flatten_with_argument_behaves_correctly
+    assert_equal [0, 1, [2], [[3]], [[[4]]]], [0, [1], [[2]], [[[3]]], [[[[4]]]]].flatten(1)
+    assert_equal [0, 1, 2, [3], [[4]]], [0, [1], [[2]], [[[3]]], [[[[4]]]]].flatten(2)
+    assert_equal [0, 1, 2, 3, [4]], [0, [1], [[2]], [[[3]]], [[[[4]]]]].flatten(3)
+    assert_equal [0, 1, 2, 3, 4], [0, [1], [[2]], [[[3]]], [[[[4]]]]].flatten(4)
+    assert_equal [Some(Some(1)), Some(Some(None)), [None]], [Some(Some(1)), Some(Some(None)), [None]].flatten(0)
+    assert_equal [Some(1), Some(None), None], [Some(Some(1)), Some(Some(None)), [None]].flatten(1)
+    assert_equal [1, None], [Some(Some(1)), Some(Some(None)), [None]].flatten(2)
+    assert_equal [1], [Some(Some(1)), Some(Some(None)), [None]].flatten(3)
   end
 end
\ No newline at end of file