[Feat#2] Nested singleton handling

lib/decors.rb

@@ -1,6 +1,6 @@
 require 'decors/decorator_base'
 require 'decors/decorator_definition'
-require 'decors/method_added_listener'
+require 'decors/method_added'
 
 module Decors
     VERSION = '0.1.0'

lib/decors/decorator_definition.rb

@@ -8,9 +8,14 @@ def define_decorator(decorator_name, decorator_class, mixin: false)
             method_definer = mixin ? :define_method : :define_singleton_method
 
             send(method_definer, decorator_name) do |*params, &blk|
-                ctx = self.singleton_class? ? ObjectSpace.each_object(self).first : self
-                ctx.send(:extend, MethodAddedListener)
-                ctx.declared_decorators << [decorator_class, params, blk]
+                if singleton_class?
+                    ObjectSpace.each_object(self).first.send(:extend, ::Decors::MethodAdded::ForwardToSingletonListener)
+                    extend(::Decors::MethodAdded::SingletonListener)
+                else
+                    extend(::Decors::MethodAdded::StandardListener)
+                end
+
+                declared_decorators << [decorator_class, params, blk]
             end
         end
     end

lib/decors/method_added.rb

@@ -0,0 +1,56 @@
+module Decors
+    module MethodAdded
+        module Handler
+            private
+
+            def declared_decorators
+                @declared_decorators ||= []
+            end
+
+            # method addition handling is the same for singleton and instance method
+            def handle_method_addition(clazz, method_name)
+                # @_ignore_additions allows to temporarily disable the hook
+                return if @_ignore_additions || declared_decorators.empty?
+                decorator_class, params, blk = declared_decorators.pop
+
+                decorated_method = clazz.instance_method(method_name)
+
+                @_ignore_additions = true
+                decorator = decorator_class.new(clazz, decorated_method, *params, &blk)
+                @_ignore_additions = false
+
+                clazz.send(:define_method, method_name) { |*args, &block| decorator.call(self, *args, &block) }
+            end
+        end
+
+        module StandardListener
+            include ::Decors::MethodAdded::Handler
+
+            def method_added(meth)
+                super
+                handle_method_addition(self, meth)
+            end
+
+            def singleton_method_added(meth)
+                super
+                handle_method_addition(singleton_class, meth)
+            end
+        end
+
+        module ForwardToSingletonListener
+            def singleton_method_added(meth)
+                super
+                singleton_class.send(:handle_method_addition, singleton_class, meth)
+            end
+        end
+
+        module SingletonListener
+            include ::Decors::MethodAdded::Handler
+
+            def singleton_method_added(meth)
+                super
+                handle_method_addition(singleton_class, meth)
+            end
+        end
+    end
+end

spec/decors_spec.rb

@@ -310,6 +310,48 @@ def self.test_method
             after { TestClass.test_method }
         end
 
+        context 'when mixin extended on the class (singleton method in singleton class)' do
+            before {
+                TestClass.class_eval {
+                    class << self
+                        extend Decors::DecoratorDefinition
+
+                        define_decorator :Deco, Deco
+
+                        Deco()
+                        def self.test_method
+                            :ok
+                        end
+                    end
+                }
+            }
+
+            it { expect(Spy).to receive(:called) }
+            after { TestClass.singleton_class.test_method }
+        end
+
+        context 'when mixin extended on the class (method in singleton class of singleton class)' do
+            before {
+                TestClass.class_eval {
+                    class << self
+                        class << self
+                            extend Decors::DecoratorDefinition
+
+                            define_decorator :Deco, Deco
+
+                            Deco()
+                            def test_method
+                                :ok
+                            end
+                        end
+                    end
+                }
+            }
+
+            it { expect(Spy).to receive(:called) }
+            after { TestClass.singleton_class.test_method }
+        end
+
         context 'when mixin extended on the class (method in singleton class)' do
             before {
                 TestClass.class_eval {
@@ -366,4 +408,3 @@ def untest_method__in_singleton
         end
     end
 end
-