1 files changed, 51 insertions, 2 deletions
diff --git a/mrbgems/mruby-enumerator/test/enumerator.rb b/mrbgems/mruby-enumerator/test/enumerator.rb
index d609cadb5..ecd6c4d65 100644
--- a/mrbgems/mruby-enumerator/test/enumerator.rb
+++ b/mrbgems/mruby-enumerator/test/enumerator.rb
@@ -6,6 +6,17 @@ class << @obj
   end
 end
 
+def assert_take(exp, enumerator)
+  result = []
+  n = exp.size
+  enumerator.each do |v|
+    result << v
+    n -= 1
+    break if n == 0
+  end if n > 0
+  assert_equal exp, result
+end
+
 assert 'Enumerator.class' do
   assert_equal Class, Enumerator.class
 end
@@ -19,7 +30,7 @@ assert 'Enumerator.new' do
   assert_equal [:x,:y,:z], [:x,:y,:z].each.map{|i| i}.sort
   assert_equal [[:x,1],[:y,2]], {x:1, y:2}.each.map{|i| i}.sort
   assert_equal [1,2,3], @obj.to_enum(:foo, 1,2,3).to_a
-  assert_equal [1,2,3], Enumerator.new { |y| i = 0; loop { y << (i += 1) } }.take(3)
+  assert_take [1,2,3], Enumerator.new { |y| i = 0; loop { y << (i += 1) } }
   assert_raise(ArgumentError) { Enumerator.new }
 
   # examples
@@ -30,7 +41,7 @@ assert 'Enumerator.new' do
       a, b = b, a + b
     end
   end
-  assert_equal [1,1,2,3,5,8,13,21,34,55], fib.take(10)
+  assert_take [1,1,2,3,5,8,13,21,34,55], fib
 end
 
 assert 'Enumerator#initialize_copy' do
@@ -549,3 +560,41 @@ assert 'Enumerable#zip' do
 
   assert_raise(TypeError) { [1].zip(1) }
 end
+
+assert 'Enumerator.produce' do
+  assert_raise(ArgumentError) { Enumerator.produce }
+
+  # Without initial object
+  passed_args = []
+  enum = Enumerator.produce {|obj| passed_args << obj; (obj || 0).succ }
+  assert_equal Enumerator, enum.class 
+  assert_take [1, 2, 3], enum
+  assert_equal [nil, 1, 2], passed_args
+
+  # With initial object
+  passed_args = []
+  enum = Enumerator.produce(1) {|obj| passed_args << obj; obj.succ }
+  assert_take [1, 2, 3], enum
+  assert_equal [1, 2], passed_args
+
+  # Raising StopIteration
+  words = %w[The quick brown fox jumps over the lazy dog]
+  enum = Enumerator.produce { words.shift or raise StopIteration }
+  assert_equal %w[The quick brown fox jumps over the lazy dog], enum.to_a
+
+  # Raising StopIteration
+  object = [[[["abc", "def"], "ghi", "jkl"], "mno", "pqr"], "stuv", "wxyz"]
+  enum = Enumerator.produce(object) {|obj|
+    obj.respond_to?(:first) or raise StopIteration
+    obj.first
+  }
+  assert_nothing_raised {
+    assert_equal [
+      [[[["abc", "def"], "ghi", "jkl"], "mno", "pqr"], "stuv", "wxyz"],
+      [[["abc", "def"], "ghi", "jkl"], "mno", "pqr"],
+      [["abc", "def"], "ghi", "jkl"],
+      ["abc", "def"],
+      "abc",
+    ], enum.to_a
+  }
+end