2 files changed, 62 insertions, 3 deletions
diff --git a/actionpack/lib/action_dispatch/middleware/flash.rb b/actionpack/lib/action_dispatch/middleware/flash.rb
index d1545355f8..1b82b15511 100644
--- a/actionpack/lib/action_dispatch/middleware/flash.rb
+++ b/actionpack/lib/action_dispatch/middleware/flash.rb
@@ -188,9 +188,6 @@ module ActionDispatch
             @used.delete(k)
           end
         end
-
-        # clean up after keys that could have been left over by calling reject! or shift on the flash
-        @used.subtract(@used - keys)
       end
 
       # Convenience accessor for flash[:alert]
diff --git a/actionpack/test/controller/flash_hash_test.rb b/actionpack/test/controller/flash_hash_test.rb
index 9b69a2648e..e7ae7ce5f5 100644
--- a/actionpack/test/controller/flash_hash_test.rb
+++ b/actionpack/test/controller/flash_hash_test.rb
@@ -75,6 +75,7 @@ module ActionDispatch
     def test_discard_no_args
       @hash['hello'] = 'world'
       @hash.discard
+
       @hash.sweep
       assert_equal({}, @hash.to_hash)
     end
@@ -83,8 +84,69 @@ module ActionDispatch
       @hash['hello'] = 'world'
       @hash['omg']   = 'world'
       @hash.discard 'hello'
+
       @hash.sweep
       assert_equal({'omg' => 'world'}, @hash.to_hash)
     end
+
+    def test_keep_sweep
+      @hash['hello'] = 'world'
+
+      @hash.sweep
+      assert_equal({'hello' => 'world'}, @hash.to_hash)
+    end
+
+    def test_update_sweep
+      @hash['hello'] = 'world'
+      @hash.update({'hi' => 'mom'})
+
+      @hash.sweep
+      assert_equal({'hello' => 'world', 'hi' => 'mom'}, @hash.to_hash)
+    end
+
+    def test_delete_sweep
+      @hash['hello'] = 'world'
+      @hash['hi']    = 'mom'
+      @hash.delete 'hi'
+
+      @hash.sweep
+      assert_equal({'hello' => 'world'}, @hash.to_hash)
+    end
+
+    def test_clear_sweep
+      @hash['hello'] = 'world'
+      @hash.clear
+
+      @hash.sweep
+      assert_equal({}, @hash.to_hash)
+    end
+
+    def test_replace_sweep
+      @hash['hello'] = 'world'
+      @hash.replace({'hi' => 'mom'})
+
+      @hash.sweep
+      assert_equal({'hi' => 'mom'}, @hash.to_hash)
+    end
+
+    def test_discard_then_add
+      @hash['hello'] = 'world'
+      @hash['omg']   = 'world'
+      @hash.discard 'hello'
+      @hash['hello'] = 'world'
+
+      @hash.sweep
+      assert_equal({'omg' => 'world', 'hello' => 'world'}, @hash.to_hash)
+    end
+
+    def test_keep_all_sweep
+      @hash['hello'] = 'world'
+      @hash['omg']   = 'world'
+      @hash.discard 'hello'
+      @hash.keep
+
+      @hash.sweep
+      assert_equal({'omg' => 'world', 'hello' => 'world'}, @hash.to_hash)
+    end
   end
 end