lib/consul/power.rb in consul-0.13.1 vs lib/consul/power.rb in consul-0.13.2

@@ -20,10 +20,11 @@
       # loaded in Consul tests.
       result.respond_to?(:load_target, true) || !!result
     end
 
     def default_include_object?(power_name, *args)
+      check_number_of_arguments_in_include_object(power_name, args.length)
       object = args.pop
       context = args
       power_value = send(power_name, *context)
       if power_value.nil?
         false
@@ -61,10 +62,20 @@
 
     def singularize_power_name(name)
       self.class.singularize_power_name(name)
     end
 
+    def check_number_of_arguments_in_include_object(power_name, given_arguments)
+      # check unmemoized methods as Memoizer wraps methods and masks the arity.
+      unmemoized_power_name = respond_to?("_unmemoized_#{power_name}") ? "_unmemoized_#{power_name}" : power_name
+      power_arity = method(unmemoized_power_name).arity
+      expected_arity = power_arity + 1 # one additional argument for the context
+      if power_arity >= 0 && expected_arity != given_arguments
+        raise ArgumentError.new("wrong number of arguments (given #{given_arguments}, expected #{expected_arity})")
+      end
+    end
+
     module ClassMethods
       include Consul::Power::DynamicAccess::ClassMethods
 
       def power(*names, &block)
         names.each do |name|
@@ -101,16 +112,27 @@
 
       def without_power(&block)
         with_power(nil, &block)
       end
 
-      def define_query_and_bang_methods(name, &query)
+      def define_query_and_bang_methods(name, options, &query)
+        is_plural = options.fetch(:is_plural)
         query_method = "#{name}?"
         bang_method = "#{name}!"
         define_method(query_method, &query)
         memoize query_method
-        define_method(bang_method) { |*args| send(query_method, *args) or powerless!(name, *args) }
+        define_method(bang_method) do |*args|
+          if is_plural
+            if send(query_method, *args)
+              send(name, *args)
+            else
+              powerless!(name, *args)
+            end
+          else
+            send(query_method, *args) or powerless!(name, *args)
+          end
+        end
         # We don't memoize the bang method since memoizer can't memoize a thrown exception
       end
 
       def define_ids_method(name)
         ids_method = power_ids_name(name)
@@ -129,17 +151,17 @@
         name = name.to_s
         if name.ends_with?('?')
           # The developer is trying to register an optimized query method
           # for singular object queries.
           name_without_suffix = name.chop
-          define_query_and_bang_methods(name_without_suffix, &block)
+          define_query_and_bang_methods(name_without_suffix, :is_plural => false, &block)
         else
           define_main_method(name, &block)
           define_ids_method(name)
-          define_query_and_bang_methods(name) { |*args| default_include_power?(name, *args) }
+          define_query_and_bang_methods(name, :is_plural => true) { |*args| default_include_power?(name, *args) }
           begin
             singular = singularize_power_name(name)
-            define_query_and_bang_methods(singular) { |*args| default_include_object?(name, *args) }
+            define_query_and_bang_methods(singular, :is_plural => false) { |*args| default_include_object?(name, *args) }
           rescue Consul::PowerNotSingularizable
             # We do not define singularized power methods if it would
             # override the collection method
           end
         end