fix content type of error response, refactor format handling

Author: Jan Xie

lib/grape/middleware/base.rb

@@ -2,7 +2,7 @@ module Grape
   module Middleware
     class Base
       attr_reader :app, :env, :options
-      
+
       # @param [Rack Application] app The standard argument for a Rack middleware.
       # @param [Hash] options A hash of options, simply stored for use by subclasses.
       def initialize(app, options = {})
@@ -38,6 +38,52 @@ def request
       def response
         Rack::Response.new(@app_response)
       end
+
+
+      module Formats
+
+        CONTENT_TYPES = {
+          :xml => 'application/xml',
+          :json => 'application/json',
+          :atom => 'application/atom+xml',
+          :rss => 'application/rss+xml',
+          :txt => 'text/plain'
+        }
+        FORMATTERS = {
+          :json => :encode_json,
+          :txt => :encode_txt,
+        }
+
+        def formatters
+          FORMATTERS.merge(options[:formatters] || {})
+        end
+
+        def content_types
+          CONTENT_TYPES.merge(options[:content_types] || {})
+        end
+
+        def content_type
+          content_types[options[:format]] || 'text/html'
+        end
+
+        def mime_types
+          content_types.invert
+        end
+
+        def formatter_for(api_format)
+          spec = formatters[api_format]
+          case spec
+          when nil
+            lambda { |obj| obj }
+          when Symbol
+            method(spec)
+          else
+            spec
+          end
+        end
+
+      end
+
     end
   end
-end
+end

lib/grape/middleware/error.rb

@@ -4,7 +4,8 @@
 module Grape
   module Middleware
     class Error < Base
-    
+      include Formats
+
       def default_options
       { 
         :default_status => 403, # default status returned on error
@@ -18,36 +19,15 @@ def default_options
       }
       end
 
-      FORMATTERS = {
-        :json => :format_json,
-        :txt => :format_txt,
-      }
-
-      def formatters
-        FORMATTERS.merge(options[:formatters])
-      end
-
-      def formatter_for(api_format)
-        spec = formatters[api_format]
-        case spec
-        when nil
-          lambda { |obj| obj }
-        when Symbol
-          method(spec)
-        else
-          spec
-        end
-      end
-
-      def format_json(message, backtrace)
+      def encode_json(message, backtrace)
         result = message.is_a?(Hash) ? message : { :error => message }
         if (options[:rescue_options] || {})[:backtrace] && backtrace && ! backtrace.empty?
           result = result.merge({ :backtrace => backtrace })
         end
         MultiJson.encode(result)
       end
 
-      def format_txt(message, backtrace)
+      def encode_txt(message, backtrace)
         result = message.is_a?(Hash) ? MultiJson.encode(message) : message
         if (options[:rescue_options] || {})[:backtrace] && backtrace && ! backtrace.empty?
           result += "\r\n "
@@ -73,7 +53,8 @@ def call!(env)
       def error_response(error = {})
         status = error[:status] || options[:default_status]
         message = error[:message] || options[:default_message]
-        headers = error[:headers] || {}
+        headers = {'Content-Type' => content_type}
+        headers.merge!(error[:headers]) if error[:headers].is_a?(Hash)
         backtrace = error[:backtrace] || []
         Rack::Response.new([format_message(message, backtrace, status)], status, headers).finish
       end

lib/grape/middleware/formatter.rb

@@ -4,18 +4,8 @@
 module Grape
   module Middleware
     class Formatter < Base
-      CONTENT_TYPES = {
-        :xml => 'application/xml',
-        :json => 'application/json',
-        :atom => 'application/atom+xml',
-        :rss => 'application/rss+xml',
-        :txt => 'text/plain'
-      }
-      FORMATTERS = {
-        :json => :encode_json,
-        :txt => :encode_txt,
-      }
-      
+      include Formats
+
       def default_options
         { 
           :default_format => :txt,
@@ -24,18 +14,6 @@ def default_options
         }
       end
 
-      def content_types
-        CONTENT_TYPES.merge(options[:content_types])
-      end
-
-      def formatters
-        FORMATTERS.merge(options[:formatters])
-      end
-      
-      def mime_types
-        content_types.invert
-      end
-      
       def headers
         env.dup.inject({}){|h,(k,v)| h[k.downcase] = v; h}
       end
@@ -92,18 +70,6 @@ def after
         Rack::Response.new(bodymap, status, headers).to_a
       end
 
-      def formatter_for(api_format)
-        spec = formatters[api_format]
-        case spec
-        when nil
-          lambda { |obj| obj }
-        when Symbol
-          method(spec)
-        else
-          spec
-        end
-      end
-      
       def encode_json(object)
         if object.respond_to? :serializable_hash
           MultiJson.encode(object.serializable_hash)

spec/grape/api_spec.rb

@@ -276,6 +276,35 @@ def app; subject end
       last_response.body.should eql 'Created'
     end
   end
+
+  context 'format' do
+    before do
+      subject.get("/foo") { "bar" }
+    end
+
+    it 'should set content type for txt format' do
+      get '/foo'
+      last_response.headers['Content-Type'].should eql 'text/plain'
+    end
+
+    it 'should set content type for json' do
+      get '/foo.json'
+      last_response.headers['Content-Type'].should eql 'application/json'
+    end
+
+    it 'should set content type for error' do
+      subject.get('/error') { error!('error in plain text', 500) }
+      get '/error'
+      last_response.headers['Content-Type'].should eql 'text/plain'
+    end
+
+    it 'should set content type for error' do
+      subject.error_format :json
+      subject.get('/error') { error!('error in json', 500) }
+      get '/error.json'
+      last_response.headers['Content-Type'].should eql 'application/json'
+    end
+  end
 
   context 'custom middleware' do
     class PhonyMiddleware