From 238a4253bf229377b686bfcecc63dda2b59cff8f Mon Sep 17 00:00:00 2001
From: kennyj <kennyj@gmail.com>
Date: Sat, 7 Jul 2012 03:34:56 +0900
Subject: [PATCH] Added support add_flash_types

---
 actionpack/CHANGELOG.md                       | 10 ++++++
 .../lib/action_controller/metal/flash.rb      | 31 ++++++++++++++-----
 actionpack/test/controller/flash_test.rb      | 26 ++++++++++++++++
 actionpack/test/template/test_case_test.rb    | 12 +++----
 4 files changed, 65 insertions(+), 14 deletions(-)

diff --git a/actionpack/CHANGELOG.md b/actionpack/CHANGELOG.md
index 7edba84ff6..67c9e04ab2 100644
--- a/actionpack/CHANGELOG.md
+++ b/actionpack/CHANGELOG.md
@@ -1,5 +1,15 @@
 ## Rails 4.0.0 (unreleased) ##
 
+*   Add `ActionController::Flash.add_flash_types` method to allow people to register their own flash types. e.g.:
+
+        class ApplicationController
+          add_flash_types :error, :warning
+        end
+
+    If you add the above code, you can use `<%= error %>` in an erb, and `rediect_to /foo, :error => 'message'` in a controller.
+
+    *kennyj*
+
 *   Remove Active Model dependency from Action Pack. *Guillermo Iguaran*
 
 *   Support unicode characters in routes. Route will be automatically escaped, so instead of manually escaping:
diff --git a/actionpack/lib/action_controller/metal/flash.rb b/actionpack/lib/action_controller/metal/flash.rb
index bd768b634e..47e71371e3 100644
--- a/actionpack/lib/action_controller/metal/flash.rb
+++ b/actionpack/lib/action_controller/metal/flash.rb
@@ -3,19 +3,34 @@ module ActionController #:nodoc:
     extend ActiveSupport::Concern
 
     included do
+      class_attribute :_flash_types, :instance_methods => false
+      self._flash_types = []
+
       delegate :flash, :to => :request
-      delegate :alert, :notice, :to => "request.flash"
-      helper_method :alert, :notice
+      add_flash_types(:alert, :notice)
+    end
+
+    module ClassMethods
+      def add_flash_types(*types)
+        types.each do |type|
+          next if _flash_types.include?(type)
+          
+          define_method(type) do
+            request.flash[type]
+          end
+          helper_method type
+
+          _flash_types << type
+        end
+      end
     end
 
     protected
       def redirect_to(options = {}, response_status_and_flash = {}) #:doc:
-        if alert = response_status_and_flash.delete(:alert)
-          flash[:alert] = alert
-        end
-
-        if notice = response_status_and_flash.delete(:notice)
-          flash[:notice] = notice
+        self.class._flash_types.each do |flash_type|
+          if type = response_status_and_flash.delete(flash_type)
+            flash[flash_type] = type
+          end
         end
 
         if other_flashes = response_status_and_flash.delete(:flash)
diff --git a/actionpack/test/controller/flash_test.rb b/actionpack/test/controller/flash_test.rb
index 32ad48b289..f91aa246f4 100644
--- a/actionpack/test/controller/flash_test.rb
+++ b/actionpack/test/controller/flash_test.rb
@@ -90,6 +90,10 @@ class FlashTest < ActionController::TestCase
     def redirect_with_other_flashes
       redirect_to '/wonderland', :flash => { :joyride => "Horses!" }
     end
+
+    def redirect_with_foo_flash
+      redirect_to "/wonderland", :foo => 'for great justice'
+    end
   end
 
   tests TestController
@@ -203,6 +207,12 @@ class FlashTest < ActionController::TestCase
     get :redirect_with_other_flashes
     assert_equal "Horses!", @controller.send(:flash)[:joyride]
   end
+
+  def test_redirect_to_with_adding_flash_types
+    @controller.class.add_flash_types :foo
+    get :redirect_with_foo_flash
+    assert_equal "for great justice", @controller.send(:flash)[:foo]
+  end
 end
 
 class FlashIntegrationTest < ActionDispatch::IntegrationTest
@@ -210,6 +220,9 @@ class FlashIntegrationTest < ActionDispatch::IntegrationTest
   SessionSecret = 'b3c631c314c0bbca50c1b2843150fe33'
 
   class TestController < ActionController::Base
+
+    add_flash_types :bar
+
     def set_flash
       flash["that"] = "hello"
       head :ok
@@ -223,6 +236,11 @@ class FlashIntegrationTest < ActionDispatch::IntegrationTest
     def use_flash
       render :inline => "flash: #{flash["that"]}"
     end
+
+    def set_bar
+      flash[:bar] = "for great justice"
+      head :ok
+    end
   end
 
   def test_flash
@@ -262,6 +280,14 @@ class FlashIntegrationTest < ActionDispatch::IntegrationTest
     end
   end
 
+  def test_added_flash_types_method
+    with_test_route_set do
+      get '/set_bar'
+      assert_response :success
+      assert_equal 'for great justice', @controller.bar
+    end
+  end
+
   private
 
     # Overwrite get to send SessionSecret in env hash
diff --git a/actionpack/test/template/test_case_test.rb b/actionpack/test/template/test_case_test.rb
index c005f040eb..387aafebd4 100644
--- a/actionpack/test/template/test_case_test.rb
+++ b/actionpack/test/template/test_case_test.rb
@@ -68,14 +68,14 @@ module ActionView
       assert_nil self.class.determine_default_helper_class("String")
     end
 
-    test "delegates notice to request.flash" do
-      view.request.flash.expects(:notice).with("this message")
-      view.notice("this message")
+    test "delegates notice to request.flash[:notice]" do
+      view.request.flash.expects(:[]).with(:notice)
+      view.notice
     end
 
-    test "delegates alert to request.flash" do
-      view.request.flash.expects(:alert).with("this message")
-      view.alert("this message")
+    test "delegates alert to request.flash[:alert]" do
+      view.request.flash.expects(:[]).with(:alert)
+      view.alert
     end
 
     test "uses controller lookup context" do