mutations on the underlying hash should also mutate the discard set

2 files changed, 24 insertions, 3 deletions

diff --git a/actionpack/lib/action_dispatch/middleware/flash.rb b/actionpack/lib/action_dispatch/middleware/flash.rb
index bf381af9e7..6564a5b16f 100644
--- a/actionpack/lib/action_dispatch/middleware/flash.rb
+++ b/actionpack/lib/action_dispatch/middleware/flash.rb
@@ -93,7 +93,7 @@ module ActionDispatch
       end
 
       def []=(k, v) #:nodoc:
-        keep(k)
+        @discard.delete k
         @flashes[k] = v
       end
 
@@ -102,7 +102,7 @@ module ActionDispatch
       end
 
       def update(h) #:nodoc:
-        h.keys.each { |k| keep(k) }
+        @discard.subtract h.keys
         @flashes.update h
         self
       end
@@ -116,6 +116,7 @@ module ActionDispatch
       end
 
       def delete(key)
+        @discard.delete key
         @flashes.delete key
         self
       end
@@ -129,6 +130,7 @@ module ActionDispatch
       end
 
       def clear
+        @discard.clear
         @flashes.clear
       end
 
@@ -139,7 +141,7 @@ module ActionDispatch
       alias :merge! :update
 
       def replace(h) #:nodoc:
-        @discard = Set.new
+        @discard.clear
         @flashes.replace h
         self
       end
diff --git a/actionpack/test/controller/flash_hash_test.rb b/actionpack/test/controller/flash_hash_test.rb
index e7ae7ce5f5..ccca0dac17 100644
--- a/actionpack/test/controller/flash_hash_test.rb
+++ b/actionpack/test/controller/flash_hash_test.rb
@@ -104,6 +104,15 @@ module ActionDispatch
       assert_equal({'hello' => 'world', 'hi' => 'mom'}, @hash.to_hash)
     end
 
+    def test_update_delete_sweep
+      @hash['hello'] = 'world'
+      @hash.delete 'hello'
+      @hash.update({'hello' => 'mom'})
+
+      @hash.sweep
+      assert_equal({'hello' => 'mom'}, @hash.to_hash)
+    end
+
     def test_delete_sweep
       @hash['hello'] = 'world'
       @hash['hi']    = 'mom'
@@ -148,5 +157,15 @@ module ActionDispatch
       @hash.sweep
       assert_equal({'omg' => 'world', 'hello' => 'world'}, @hash.to_hash)
     end
+
+    def test_double_sweep
+      @hash['hello'] = 'world'
+      @hash.sweep
+
+      assert_equal({'hello' => 'world'}, @hash.to_hash)
+
+      @hash.sweep
+      assert_equal({}, @hash.to_hash)
+    end
   end
 end