1 files changed, 51 insertions, 2 deletions
diff --git a/spec/ruby/core/enumerable/flat_map_spec.rb b/spec/ruby/core/enumerable/flat_map_spec.rb
index bd07eab6c5..ef50cb2696 100644
--- a/spec/ruby/core/enumerable/flat_map_spec.rb
+++ b/spec/ruby/core/enumerable/flat_map_spec.rb
@@ -1,7 +1,56 @@
 require_relative '../../spec_helper'
 require_relative 'fixtures/classes'
-require_relative 'shared/collect_concat'
+require_relative 'shared/enumerable_enumeratorized'
 
 describe "Enumerable#flat_map" do
-  it_behaves_like :enumerable_collect_concat, :flat_map
+  it "yields elements to the block and flattens one level" do
+    numerous = EnumerableSpecs::Numerous.new(1, [2, 3], [4, [5, 6]], {foo: :bar})
+    numerous.flat_map { |i| i }.should == [1, 2, 3, 4, [5, 6], {foo: :bar}]
+  end
+
+  it "appends non-Array elements that do not define #to_ary" do
+    obj = mock("to_ary undefined")
+
+    numerous = EnumerableSpecs::Numerous.new(1, obj, 2)
+    numerous.flat_map { |i| i }.should == [1, obj, 2]
+  end
+
+  it "concatenates the result of calling #to_ary if it returns an Array" do
+    obj = mock("to_ary defined")
+    obj.should_receive(:to_ary).and_return([:a, :b])
+
+    numerous = EnumerableSpecs::Numerous.new(1, obj, 2)
+    numerous.flat_map { |i| i }.should == [1, :a, :b, 2]
+  end
+
+  it "does not call #to_a" do
+    obj = mock("to_ary undefined")
+    obj.should_not_receive(:to_a)
+
+    numerous = EnumerableSpecs::Numerous.new(1, obj, 2)
+    numerous.flat_map { |i| i }.should == [1, obj, 2]
+  end
+
+  it "appends an element that defines #to_ary that returns nil" do
+    obj = mock("to_ary defined")
+    obj.should_receive(:to_ary).and_return(nil)
+
+    numerous = EnumerableSpecs::Numerous.new(1, obj, 2)
+    numerous.flat_map { |i| i }.should == [1, obj, 2]
+  end
+
+  it "raises a TypeError if an element defining #to_ary does not return an Array or nil"  do
+    obj = mock("to_ary defined")
+    obj.should_receive(:to_ary).and_return("array")
+
+    -> { [1, obj, 3].flat_map { |i| i } }.should.raise(TypeError)
+  end
+
+  it "returns an enumerator when no block given" do
+    enum = EnumerableSpecs::Numerous.new(1, 2).flat_map
+    enum.should.instance_of?(Enumerator)
+    enum.each{ |i| [i] * i }.should == [1, 2, 2]
+  end
+
+  it_behaves_like :enumerable_enumeratorized_with_origin_size, :flat_map
 end