Merge pull request #1125 from yannvanhalewyn/improve-fetch-on-session-hash

Improve and test SessionHash#fetch

2 files changed, 28 insertions, 1 deletions

diff --git a/lib/rack/session/abstract/id.rb b/lib/rack/session/abstract/id.rb
index 4c139637..ca1a2628 100644
--- a/lib/rack/session/abstract/id.rb
+++ b/lib/rack/session/abstract/id.rb
@@ -18,6 +18,8 @@ module Rack
         include Enumerable
         attr_writer :id
 
+        Unspecified = Object.new
+
         def self.find(req)
           req.get_header RACK_SESSION
         end
@@ -54,7 +56,15 @@ module Rack
           load_for_read!
           @data[key.to_s]
         end
-        alias :fetch :[]
+
+        def fetch(key, default=Unspecified, &block)
+          load_for_read!
+          if default == Unspecified
+            @data.fetch(key.to_s, &block)
+          else
+            @data.fetch(key.to_s, default, &block)
+          end
+        end
 
         def has_key?(key)
           load_for_read!
diff --git a/test/spec_session_abstract_session_hash.rb b/test/spec_session_abstract_session_hash.rb
index 6d73a80a..76b34a01 100644
--- a/test/spec_session_abstract_session_hash.rb
+++ b/test/spec_session_abstract_session_hash.rb
@@ -25,4 +25,21 @@ describe Rack::Session::Abstract::SessionHash do
     assert_equal [:bar, :qux], hash.values
   end
 
+  describe "#fetch" do
+    it "returns value for a matching key" do
+      assert_equal :bar, hash.fetch(:foo)
+    end
+
+    it "works with a default value" do
+      assert_equal :default, hash.fetch(:unknown, :default)
+    end
+
+    it "works with a block" do
+      assert_equal :default, hash.fetch(:unkown) { :default }
+    end
+
+    it "it raises when fetching unknown keys without defaults" do
+      lambda { hash.fetch(:unknown) }.must_raise KeyError
+    end
+  end
 end