1 files changed, 25 insertions, 4 deletions
diff --git a/actionpack/test/controller/flash_test.rb b/actionpack/test/controller/flash_test.rb
index 37ccc7a4a5..8340aab4d2 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,9 +220,7 @@ class FlashIntegrationTest < ActionDispatch::IntegrationTest
   SessionSecret = 'b3c631c314c0bbca50c1b2843150fe33'
 
   class TestController < ActionController::Base
-    def dont_set_flash
-      head :ok
-    end
+    add_flash_types :bar
 
     def set_flash
       flash["that"] = "hello"
@@ -227,6 +235,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
@@ -266,6 +279,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
@@ -277,7 +298,7 @@ class FlashIntegrationTest < ActionDispatch::IntegrationTest
     def with_test_route_set
       with_routing do |set|
         set.draw do
-          match ':action', :to => FlashIntegrationTest::TestController
+          get ':action', :to => FlashIntegrationTest::TestController
         end
 
         @app = self.class.build_app(set) do |middleware|