configurator: allow hooks to be passed callable objects

Premade lambda/proc/Proc objects may all be passed, to the
hooks, not just anonymous blocks.

2 files changed, 21 insertions, 6 deletions

diff --git a/lib/unicorn/configurator.rb b/lib/unicorn/configurator.rb
index 2dcad48..150377b 100644
--- a/lib/unicorn/configurator.rb
+++ b/lib/unicorn/configurator.rb
@@ -105,15 +105,15 @@ module Unicorn
     #    # HUP after QUIT-ing the original master:
     #    server.listen("127.0.0.1:#{9293 + worker_nr}") rescue nil
     #  end
-    def after_fork(&block)
-      set_hook(:after_fork, block)
+    def after_fork(*args, &block)
+      set_hook(:after_fork, block_given? ? block : args[0])
     end
 
     # sets before_fork got be a given Proc object.  This Proc
     # object will be called by the master process before forking
     # each worker.
-    def before_fork(&block)
-      set_hook(:before_fork, block)
+    def before_fork(*args, &block)
+      set_hook(:before_fork, block_given? ? block : args[0])
     end
 
     # sets the before_exec hook to a given Proc object.  This
@@ -122,8 +122,8 @@ module Unicorn
     # for freeing certain OS resources that you do NOT wish to
     # share with the reexeced child process.
     # There is no corresponding after_exec hook (for obvious reasons).
-    def before_exec(&block)
-      set_hook(:before_exec, block, 1)
+    def before_exec(*args, &block)
+      set_hook(:before_exec, block_given? ? block : args[0], 1)
     end
 
     # sets the timeout of worker processes to +seconds+.  Workers
diff --git a/test/unit/test_configurator.rb b/test/unit/test_configurator.rb
index 284d727..46fc248 100644
--- a/test/unit/test_configurator.rb
+++ b/test/unit/test_configurator.rb
@@ -91,4 +91,19 @@ class TestConfigurator < Test::Unit::TestCase
     end
   end
 
+  def test_after_fork_proc
+    [ proc { |a,b| }, Proc.new { |a,b| }, lambda { |a,b| } ].each do |my_proc|
+      Unicorn::Configurator.new(:after_fork => my_proc).commit!(self)
+      assert_equal my_proc, @after_fork
+    end
+  end
+
+  def test_after_fork_wrong_arity
+    [ proc { |a| }, Proc.new { }, lambda { |a,b,c| } ].each do |my_proc|
+      assert_raises(ArgumentError) do
+        Unicorn::Configurator.new(:after_fork => my_proc)
+      end
+    end
+  end
+
 end