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

Pre-discard flash messages

4 files changed, 51 insertions, 28 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:
diff --git a/actionpack/test/controller/flash_hash_test.rb b/actionpack/test/controller/flash_hash_test.rb
index d979b561f2..081288ef21 100644
--- a/actionpack/test/controller/flash_hash_test.rb
+++ b/actionpack/test/controller/flash_hash_test.rb
@@ -57,33 +57,36 @@ module ActionDispatch
 
     def test_to_session_value
       @hash['foo'] = 'bar'
-      assert_equal({'flashes' => {'foo' => 'bar'}, 'discard' => []}, @hash.to_session_value)
-
-      @hash.discard('foo')
-      assert_equal({'flashes' => {'foo' => 'bar'}, 'discard' => %w[foo]}, @hash.to_session_value)
+      assert_equal({'flashes' => {'foo' => 'bar'}}, @hash.to_session_value)
 
       @hash.now['qux'] = 1
-      assert_equal({'flashes' => {'foo' => 'bar', 'qux' => 1}, 'discard' => %w[foo qux]}, @hash.to_session_value)
+      assert_equal({'flashes' => {'foo' => 'bar'}}, @hash.to_session_value)
+
+      @hash.discard('foo')
+      assert_equal(nil, @hash.to_session_value)
 
       @hash.sweep
       assert_equal(nil, @hash.to_session_value)
     end
 
     def test_from_session_value
-      rails_3_2_cookie = 'BAh7B0kiD3Nlc3Npb25faWQGOgZFRkkiJWY4ZTFiODE1MmJhNzYwOWMyOGJiYjE3ZWM5MjYzYmE3BjsAVEkiCmZsYXNoBjsARm86JUFjdGlvbkRpc3BhdGNoOjpGbGFzaDo6Rmxhc2hIYXNoCToKQHVzZWRvOghTZXQGOgpAaGFzaHsAOgxAY2xvc2VkRjoNQGZsYXNoZXN7BkkiDG1lc3NhZ2UGOwBGSSIKSGVsbG8GOwBGOglAbm93MA=='
+      # {"session_id"=>"f8e1b8152ba7609c28bbb17ec9263ba7", "flash"=>#<ActionDispatch::Flash::FlashHash:0x00000000000000 @used=#<Set: {"farewell"}>, @closed=false, @flashes={"greeting"=>"Hello", "farewell"=>"Goodbye"}, @now=nil>}
+      rails_3_2_cookie = 'BAh7B0kiD3Nlc3Npb25faWQGOgZFRkkiJWY4ZTFiODE1MmJhNzYwOWMyOGJiYjE3ZWM5MjYzYmE3BjsAVEkiCmZsYXNoBjsARm86JUFjdGlvbkRpc3BhdGNoOjpGbGFzaDo6Rmxhc2hIYXNoCToKQHVzZWRvOghTZXQGOgpAaGFzaHsGSSINZmFyZXdlbGwGOwBUVDoMQGNsb3NlZEY6DUBmbGFzaGVzewdJIg1ncmVldGluZwY7AFRJIgpIZWxsbwY7AFRJIg1mYXJld2VsbAY7AFRJIgxHb29kYnllBjsAVDoJQG5vdzA='
       session = Marshal.load(Base64.decode64(rails_3_2_cookie))
       hash = Flash::FlashHash.from_session_value(session['flash'])
-      assert_equal({'flashes' => {'message' => 'Hello'}, 'discard' => %w[message]}, hash.to_session_value)
+      assert_equal({'greeting' => 'Hello'}, hash.to_hash)
+      assert_equal(nil, hash.to_session_value)
     end
 
     def test_from_session_value_on_json_serializer
-      decrypted_data = "{ \"session_id\":\"d98bdf6d129618fc2548c354c161cfb5\", \"flash\":{\"discard\":[], \"flashes\":{\"message\":\"hey you\"}} }"
+      decrypted_data = "{ \"session_id\":\"d98bdf6d129618fc2548c354c161cfb5\", \"flash\":{\"discard\":[\"farewell\"], \"flashes\":{\"greeting\":\"Hello\",\"farewell\":\"Goodbye\"}} }"
       session = ActionDispatch::Cookies::JsonSerializer.load(decrypted_data)
       hash = Flash::FlashHash.from_session_value(session['flash'])
 
-      assert_equal({'discard' => %w[message], 'flashes' => { 'message' => 'hey you'}}, hash.to_session_value)
-      assert_equal "hey you", hash[:message]
-      assert_equal "hey you", hash["message"]
+      assert_equal({'greeting' => 'Hello'}, hash.to_hash)
+      assert_equal(nil, hash.to_session_value)
+      assert_equal "Hello", hash[:greeting]
+      assert_equal "Hello", hash["greeting"]
     end
 
     def test_empty?
diff --git a/actionpack/test/controller/test_case_test.rb b/actionpack/test/controller/test_case_test.rb
index 901db63e39..25337a7dd4 100644
--- a/actionpack/test/controller/test_case_test.rb
+++ b/actionpack/test/controller/test_case_test.rb
@@ -14,6 +14,11 @@ class TestCaseTest < ActionController::TestCase
       render text: 'ignore me'
     end
 
+    def delete_flash
+      flash.delete("test")
+      render :text => 'ignore me'
+    end
+
     def set_flash_now
       flash.now["test_now"] = ">#{flash["test_now"]}<"
       render text: 'ignore me'
@@ -290,6 +295,13 @@ XML
     assert_equal '>value_now<', flash['test_now']
   end
 
+  def test_process_delete_flash
+    process :set_flash
+    process :delete_flash
+    assert_empty flash
+    assert_empty session
+  end
+
   def test_process_with_session
     process :set_session
     assert_equal 'A wonder', session['string'], "A value stored in the session should be available by string key"