Added support add_flash_types

1 files changed, 26 insertions, 0 deletions

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