2 files changed, 81 insertions, 4 deletions

diff --git a/activesupport/lib/active_support/lazy_load_hooks.rb b/activesupport/lib/active_support/lazy_load_hooks.rb
index 642a4c105c..3664431a28 100644
--- a/activesupport/lib/active_support/lazy_load_hooks.rb
+++ b/activesupport/lib/active_support/lazy_load_hooks.rb
@@ -2,16 +2,26 @@ module ActiveSupport
   @load_hooks = Hash.new {|h,k| h[k] = [] }
   @loaded = {}
 
-  def self.on_load(name, &block)
+  def self.on_load(name, options = {}, &block)
     if base = @loaded[name]
-      base.instance_eval(&block)
+      execute_hook(base, options, block)
     else
-      @load_hooks[name] << block
+      @load_hooks[name] << [block, options]
+    end
+  end
+
+  def self.execute_hook(base, options, block)
+    if options[:yield]
+      block.call(base)
+    else
+      base.instance_eval(&block)
     end
   end
 
   def self.run_load_hooks(name, base = Object)
-    @load_hooks[name].each { |hook| base.instance_eval(&hook) }
     @loaded[name] = base
+    @load_hooks[name].each do |hook, options|
+      execute_hook(base, options, hook)
+    end
   end
 end
\ No newline at end of file
diff --git a/activesupport/test/lazy_load_hooks_test.rb b/activesupport/test/lazy_load_hooks_test.rb
new file mode 100644
index 0000000000..58ccc14324
--- /dev/null
+++ b/activesupport/test/lazy_load_hooks_test.rb
@@ -0,0 +1,67 @@
+require 'abstract_unit'
+
+class LazyLoadHooksTest < ActiveSupport::TestCase
+  def test_basic_hook
+    i = 0
+    ActiveSupport.on_load(:basic_hook) { i += 1 }
+    ActiveSupport.run_load_hooks(:basic_hook)
+    assert_equal 1, i
+  end
+
+  def test_hook_registered_after_run
+    i = 0
+    ActiveSupport.run_load_hooks(:registered_after)
+    assert_equal 0, i
+    ActiveSupport.on_load(:registered_after) { i += 1 }
+    assert_equal 1, i
+  end
+
+  def test_hook_receives_a_context
+    i = 0
+    ActiveSupport.on_load(:contextual) { i += incr }
+    assert_equal 0, i
+    ActiveSupport.run_load_hooks(:contextual, FakeContext.new(2))
+    assert_equal 2, i
+  end
+
+  def test_hook_receives_a_context_afterward
+    i = 0
+    ActiveSupport.run_load_hooks(:contextual_after, FakeContext.new(2))
+    assert_equal 0, i
+    ActiveSupport.on_load(:contextual_after) { i += incr }
+    assert_equal 2, i
+  end
+
+  def test_hook_with_yield_true
+    i = 0
+    ActiveSupport.on_load(:contextual_yield, :yield => true) do |obj|
+      i += obj.incr + incr_amt
+    end
+    assert_equal 0, i
+    ActiveSupport.run_load_hooks(:contextual_yield, FakeContext.new(2))
+    assert_equal 7, i
+  end
+
+  def test_hook_with_yield_true_afterward
+    i = 0
+    ActiveSupport.run_load_hooks(:contextual_yield_after, FakeContext.new(2))
+    assert_equal 0, i
+    ActiveSupport.on_load(:contextual_yield_after, :yield => true) do |obj|
+      i += obj.incr + incr_amt
+    end
+    assert_equal 7, i
+  end
+
+private
+
+  def incr_amt
+    5
+  end
+
+  class FakeContext
+    attr_reader :incr
+    def initialize(incr)
+      @incr = incr
+    end
+  end
+end
\ No newline at end of file