From 2ffa278651704a0ed0bc954b8d3f22190a4b61fc Mon Sep 17 00:00:00 2001
From: Ryan Briones <ryan.briones@brionesandco.com>
Date: Thu, 23 Oct 2008 23:16:37 -0400
Subject: [PATCH] add ability to set custom headers; add http_accept helper

---
 lib/webrat/core/session.rb       | 34 +++++++++++++++++++++++++--
 spec/webrat/core/session_spec.rb | 40 ++++++++++++++++++++++++++++++--
 2 files changed, 70 insertions(+), 4 deletions(-)

diff --git a/lib/webrat/core/session.rb b/lib/webrat/core/session.rb
index ab89e50..b630679 100644
--- a/lib/webrat/core/session.rb
+++ b/lib/webrat/core/session.rb
@@ -13,6 +13,7 @@ module Webrat
       @http_method     = :get
       @data            = {}
       @default_headers = {}
+      @custom_headers  = {}
     end
 
     # Saves the page out to RAILS_ROOT/tmp/ and opens it in the default
@@ -52,13 +53,42 @@ module Webrat
     def saved_page_dir
       File.expand_path(".")
     end
+
+    def header(key, value)
+      @custom_headers[key] = value
+    end
+
+    def http_accept(mime_type_as_string_or_symbol)
+      if String === mime_type_as_string_or_symbol
+        header('Accept', mime_type_as_string_or_symbol)
+      elsif Symbol === mime_type_as_string_or_symbol
+        mime_type_as_string = case mime_type_as_string_or_symbol
+                              when :text then "text/plain"
+                              when :html then "text/html"
+                              when :js then "text/javascript"
+                              when :css then "text/css"
+                              when :ics then "text/calendar"
+                              when :csv then "text/csv"
+                              when :xml then "application/xml"
+                              when :rss then "application/rss+xml"
+                              when :atom then "application/atom+xml"
+                              when :yaml then "application/x-yaml"
+                              when :multipart_form then "multipart/form-data"
+                              when :url_encoded_form then "application/x-www-form-urlencoded"
+                              when :json then "application/json"
+                              end
+        raise ArgumentError.new("Invalid Mime type: #{mime_type_as_string_or_symbol.inspect}") unless mime_type_as_string
+        header('Accept', mime_type_as_string)
+      end
+    end
     
     def basic_auth(user, pass)
-      @default_headers['HTTP_AUTHORIZATION'] = "Basic " + ["#{user}:#{pass}"].pack("m*")
+      encoded_login = ["#{user}:#{pass}"].pack("m*")
+      header('HTTP_AUTHORIZATION', "Basic #{encoded_login}")
     end
 
     def headers
-      @default_headers.dup
+      @default_headers.dup.merge(@custom_headers.dup)
     end
 
     def request_page(url, http_method, data)
diff --git a/spec/webrat/core/session_spec.rb b/spec/webrat/core/session_spec.rb
index 376e4b7..de11c8b 100644
--- a/spec/webrat/core/session_spec.rb
+++ b/spec/webrat/core/session_spec.rb
@@ -19,7 +19,7 @@ describe Webrat::Session do
   
   it "should open the page in the browser" do
     session = Webrat::Session.new
-    session.should_receive(:`).with("open path")
+    session.should_receive(:`).with("open path") #`)
     session.open_in_browser("path")
   end
   
@@ -29,5 +29,41 @@ describe Webrat::Session do
     current_page.should_not be_nil
     current_page.should respond_to(:url)
   end
+
+  it "should allow custom headers to be set" do
+    session = Webrat::Session.new
+    session.header('Accept', 'application/xml')
+
+    session.instance_variable_get(:@custom_headers)['Accept'].should == 'application/xml'
+  end
+
+  it "should return a copy of the headers to be sent" do
+    session = Webrat::Session.new
+    session.instance_eval { 
+      @default_headers = {'HTTP_X_FORWARDED_FOR' => '192.168.1.1'}
+      @custom_headers = {'Accept' => 'application/xml'}
+    }
+    session.headers.should == {'HTTP_X_FORWARDED_FOR' => '192.168.1.1', 'Accept' => 'application/xml'}
+  end
+
+  describe "#http_accept" do
+    before(:each) do
+      @session = Webrat::Session.new
+    end
+
+    it "should set the Accept header with the string Mime type" do
+      @session.http_accept('application/xml')
+      @session.headers['Accept'].should == 'application/xml'
+    end
+
+    it "should set the Accept head with the string value of the symbol Mime type" do
+      @session.http_accept(:xml)
+      @session.headers['Accept'].should == 'application/xml'
+    end
+
+    it "should raise an error if a symbol Mime type is passed that does not exist" do
+      lambda { @session.http_accept(:oogabooga) }.should raise_error(ArgumentError)
+    end
+  end
   
-end
\ No newline at end of file
+end