Allow reloader to be configured.

2 files changed, 51 insertions, 14 deletions

diff --git a/actionpack/lib/action_dispatch/middleware/reloader.rb b/actionpack/lib/action_dispatch/middleware/reloader.rb
index 29289a76b4..4f48f1c974 100644
--- a/actionpack/lib/action_dispatch/middleware/reloader.rb
+++ b/actionpack/lib/action_dispatch/middleware/reloader.rb
@@ -43,34 +43,58 @@ module ActionDispatch
 
     # Execute all prepare callbacks.
     def self.prepare!
-      new(nil).run_callbacks :prepare
+      new(nil).prepare!
     end
 
     # Execute all cleanup callbacks.
     def self.cleanup!
-      new(nil).run_callbacks :cleanup
+      new(nil).cleanup!
     end
 
-    def initialize(app)
+    def initialize(app, condition=nil)
       @app = app
-    end
-
-    module CleanupOnClose
-      def close
-        super if defined?(super)
-      ensure
-        ActionDispatch::Reloader.cleanup!
-      end
+      @condition = condition || lambda { true }
+      @validated = true
     end
 
     def call(env)
-      run_callbacks :prepare
+      @validated = @condition.call
+      prepare!
       response = @app.call(env)
-      response[2].extend(CleanupOnClose)
+      response[2].extend(module_hook)
       response
     rescue Exception
-      run_callbacks :cleanup
+      cleanup!
       raise
     end
+
+    def prepare! #:nodoc:
+      run_callbacks :prepare if validated?
+    end
+
+    def cleanup! #:nodoc:
+      run_callbacks :cleanup if validated?
+    ensure
+      @validated = true
+    end
+
+    private
+
+    def validated? #:nodoc:
+      @validated
+    end
+
+    def module_hook #:nodoc:
+      middleware = self
+      Module.new do
+        define_method :close do
+          begin
+            super() if defined?(super)
+          ensure
+            middleware.cleanup!
+          end
+        end
+      end
+    end
   end
 end
diff --git a/actionpack/test/dispatch/reloader_test.rb b/actionpack/test/dispatch/reloader_test.rb
index eaabc1feb3..bd24256427 100644
--- a/actionpack/test/dispatch/reloader_test.rb
+++ b/actionpack/test/dispatch/reloader_test.rb
@@ -43,6 +43,19 @@ class ReloaderTest < Test::Unit::TestCase
     assert_respond_to body, :close
   end
 
+  def test_condition_specifies_when_to_reload
+    i, j = 0, 0, 0, 0
+    Reloader.to_prepare { |*args| i += 1 }
+    Reloader.to_cleanup { |*args| j += 1 }
+    app = Reloader.new(lambda { |env| [200, {}, []] }, lambda { i < 3 })
+    5.times do
+      resp = app.call({})
+      resp[2].close
+    end
+    assert_equal 3, i
+    assert_equal 3, j
+  end
+
   def test_returned_body_object_behaves_like_underlying_object
     body = call_and_return_body do
       b = MyBody.new