Merge pull request #163 from appsignal/instrument-method-blocks

Support blocks arguments on Method instrumentation

lib/appsignal/integrations/object.rb

@@ -2,24 +2,24 @@ class Object
   def self.appsignal_instrument_class_method(method_name, options = {})
     singleton_class.send \
       :alias_method, "appsignal_uninstrumented_#{method_name}", method_name
-    singleton_class.send(:define_method, method_name) do |*args|
+    singleton_class.send(:define_method, method_name) do |*args, &block|
       name = options.fetch(:name) do
         "#{method_name}.class_method.#{appsignal_reverse_class_name}.other"
       end
       Appsignal.instrument name do
-        send "appsignal_uninstrumented_#{method_name}", *args
+        send "appsignal_uninstrumented_#{method_name}", *args, &block
       end
     end
   end
 
   def self.appsignal_instrument_method(method_name, options = {})
     alias_method "appsignal_uninstrumented_#{method_name}", method_name
-    define_method method_name do |*args|
+    define_method method_name do |*args, &block|
       name = options.fetch(:name) do
         "#{method_name}.#{appsignal_reverse_class_name}.other"
       end
       Appsignal.instrument name do
-        send "appsignal_uninstrumented_#{method_name}", *args
+        send "appsignal_uninstrumented_#{method_name}", *args, &block
       end
     end
   end

spec/lib/appsignal/integrations/object_spec.rb

@@ -92,6 +92,21 @@ def foo
             expect(instance.foo).to eq(1)
           end
         end
+
+        context "with a method given a block" do
+          let(:klass) do
+            Class.new do
+              def foo
+                yield
+              end
+              appsignal_instrument_method :foo
+            end
+          end
+
+          it "should yield the block" do
+            expect(instance.foo { 42 }).to eq(42)
+          end
+        end
       end
 
       context "when not active" do
@@ -194,6 +209,21 @@ def self.bar
             expect(klass.bar).to eq(2)
           end
         end
+
+        context "with a method given a block" do
+          let(:klass) do
+            Class.new do
+              def self.bar
+                yield
+              end
+              appsignal_instrument_class_method :bar
+            end
+          end
+
+          it "should yield the block" do
+            expect(klass.bar { 42 }).to eq(42)
+          end
+        end
       end
 
       context "when not active" do