Merge pull request #18721 from sj26/pre-discard-flash

Pre-discard flash messages

2 files changed, 25 insertions, 17 deletions

diff --git a/actionpack/lib/action_controller/test_case.rb b/actionpack/lib/action_controller/test_case.rb
index afe95e3fec..b05700aace 100644
--- a/actionpack/lib/action_controller/test_case.rb
+++ b/actionpack/lib/action_controller/test_case.rb
@@ -677,6 +677,8 @@ module ActionController
 
         if flash_value = @request.flash.to_session_value
           @request.session['flash'] = flash_value
+        else
+          @request.session.delete('flash')
         end
 
         @response
diff --git a/actionpack/lib/action_dispatch/middleware/flash.rb b/actionpack/lib/action_dispatch/middleware/flash.rb
index a7f95150a4..59639a010e 100644
--- a/actionpack/lib/action_dispatch/middleware/flash.rb
+++ b/actionpack/lib/action_dispatch/middleware/flash.rb
@@ -80,24 +80,30 @@ module ActionDispatch
       include Enumerable
 
       def self.from_session_value(value) #:nodoc:
-        flash = case value
-                when FlashHash # Rails 3.1, 3.2
-                  new(value.instance_variable_get(:@flashes), value.instance_variable_get(:@used))
-                when Hash # Rails 4.0
-                  new(value['flashes'], value['discard'])
-                else
-                  new
-                end
-
-        flash.tap(&:sweep)
-      end
-      
-      # Builds a hash containing the discarded values and the hashes
-      # representing the flashes.
-      # If there are no values in @flashes, returns nil.
+        case value
+        when FlashHash # Rails 3.1, 3.2
+          flashes = value.instance_variable_get(:@flashes)
+          if discard = value.instance_variable_get(:@used)
+            flashes.except!(*discard)
+          end
+          new(flashes, flashes.keys)
+        when Hash # Rails 4.0
+          flashes = value['flashes']
+          if discard = value['discard']
+            flashes.except!(*discard)
+          end
+          new(flashes, flashes.keys)
+        else
+          new
+        end
+      end
+
+      # Builds a hash containing the flashes to keep for the next request.
+      # If there are none to keep, returns nil.
       def to_session_value #:nodoc:
-        return nil if empty?
-        {'discard' => @discard.to_a, 'flashes' => @flashes}
+        flashes_to_keep = @flashes.except(*@discard)
+        return nil if flashes_to_keep.empty?
+        {'flashes' => flashes_to_keep}
       end
 
       def initialize(flashes = {}, discard = []) #:nodoc: