Define `#+` to `Enumerator` and `Enumerator#Chain` instead of `Enumerable`

2 files changed, 20 insertions, 5 deletions

diff --git a/mrbgems/mruby-enum-chain/mrblib/chain.rb b/mrbgems/mruby-enum-chain/mrblib/chain.rb
index 52f5f0656..03acfb7ef 100644
--- a/mrbgems/mruby-enum-chain/mrblib/chain.rb
+++ b/mrbgems/mruby-enum-chain/mrblib/chain.rb
@@ -6,13 +6,13 @@ module Enumerable
   def chain(*args)
     Enumerator::Chain.new(self, *args)
   end
+end
 
+class Enumerator
   def +(other)
-    Enumerator::Chain.new(self, other)
+    Chain.new(self, other)
   end
-end
 
-class Enumerator
   class Chain
     include Enumerable
 
@@ -50,6 +50,10 @@ class Enumerator
       self
     end
 
+    def +(other)
+      self.class.new(self, other)
+    end
+
     def inspect
       "#<#{self.class}: #{@enums.inspect}>"
     end
diff --git a/mrbgems/mruby-enum-chain/test/enum_chain.rb b/mrbgems/mruby-enum-chain/test/enum_chain.rb
index 9aea783b9..1d3d691ca 100644
--- a/mrbgems/mruby-enum-chain/test/enum_chain.rb
+++ b/mrbgems/mruby-enum-chain/test/enum_chain.rb
@@ -12,7 +12,7 @@ assert("Enumerable#chain") do
   assert_raise(NoMethodError) { c.chain }
 end
 
-assert("Enumerable#+") do
+assert("Enumerator#+") do
   a = [].each
   b = {}.each
   c = Object.new # not has #each method
@@ -24,7 +24,7 @@ assert("Enumerable#+") do
   assert_raise(NoMethodError) { c + a }
 end
 
-assert("Enumerator.new") do
+assert("Enumerator::Chain.new") do
   a = []
   b = {}
   c = Object.new # not has #each method
@@ -84,3 +84,14 @@ assert("Enumerator::Chain#rewind") do
   c = e1.chain(e2).each{}.rewind
   assert_equal [e2.__id__, e1.__id__], rewound
 end
+
+assert("Enumerator::Chain#+") do
+  a = [].chain
+  b = {}.chain
+  c = Object.new # not has #each method
+
+  assert_kind_of Enumerator::Chain, a + b
+  assert_kind_of Enumerator::Chain, a + c
+  assert_kind_of Enumerator::Chain, b + a
+  assert_kind_of Enumerator::Chain, b + c
+end