ActionController::Flash::FlashHash.use now returns either the value corresponding to the passed key or itself when no key is passed [#1792 state:resolved]

Signed-off-by: Joshua Peek <josh@joshpeek.com>

2 files changed, 20 insertions, 6 deletions

diff --git a/actionpack/lib/action_controller/base/chained/flash.rb b/actionpack/lib/action_controller/base/chained/flash.rb
index 6bd482d85a..7a8dd2dcf9 100644
--- a/actionpack/lib/action_controller/base/chained/flash.rb
+++ b/actionpack/lib/action_controller/base/chained/flash.rb
@@ -135,12 +135,11 @@ module ActionController #:nodoc:
         #     use('msg')          # marks the "msg" entry as used
         #     use(nil, false)     # marks the entire flash as unused (keeps it around for one more action)
         #     use('msg', false)   # marks the "msg" entry as unused (keeps it around for one more action)
-        def use(k=nil, v=true)
-          unless k.nil?
-            @used[k] = v
-          else
-            keys.each{ |key| use(key, v) }
-          end
+        # Returns the single value for the key you asked to be marked (un)used or the FlashHash itself
+        # if no key is passed.
+        def use(key = nil, used = true)
+          Array(key || keys).each { |k| @used[k] = used }
+          return key ? self[key] : self
         end
     end
 
diff --git a/actionpack/test/controller/flash_test.rb b/actionpack/test/controller/flash_test.rb
index 84e27d7779..ee215bb19e 100644
--- a/actionpack/test/controller/flash_test.rb
+++ b/actionpack/test/controller/flash_test.rb
@@ -140,4 +140,19 @@ class FlashTest < ActionController::TestCase
     get :std_action
     assert_nil assigns["flash_copy"]["foo"]
   end
+
+  def test_keep_and_discard_return_values
+    flash = ActionController::Flash::FlashHash.new
+    flash.update(:foo => :foo_indeed, :bar => :bar_indeed)
+
+    assert_equal(:foo_indeed, flash.discard(:foo)) # valid key passed
+    assert_nil flash.discard(:unknown) # non existant key passed
+    assert_equal({:foo => :foo_indeed, :bar => :bar_indeed}, flash.discard()) # nothing passed
+    assert_equal({:foo => :foo_indeed, :bar => :bar_indeed}, flash.discard(nil)) # nothing passed      
+
+    assert_equal(:foo_indeed, flash.keep(:foo)) # valid key passed
+    assert_nil flash.keep(:unknown) # non existant key passed
+    assert_equal({:foo => :foo_indeed, :bar => :bar_indeed}, flash.keep()) # nothing passed
+    assert_equal({:foo => :foo_indeed, :bar => :bar_indeed}, flash.keep(nil)) # nothing passed     
+  end
 end