lib/restfulness/response.rb in restfulness-0.1.0 vs lib/restfulness/response.rb in restfulness-0.2.0

@@ -5,18 +5,15 @@
 
     # Incoming data
     attr_reader :request
 
     # Outgoing data
-    attr_reader :code, :headers, :payload
+    attr_reader :status, :headers, :payload
 
-
     def initialize(request)
       @request = request
-
-      # Default headers
-      @headers = {'Content-Type' => 'application/json; charset=utf-8'}
+      @headers = {}
     end
 
     def run
       logger.info("Responding to #{request.action.to_s.upcase} #{request.uri.to_s} from #{request.remote_ip}")
 
@@ -29,28 +26,53 @@
         resource.check_callbacks
 
         # Perform the actual work
         result = resource.call
 
-        @code    ||= (result ? 200 : 204)
-        @payload   = MultiJson.encode(result)
+        update_status_and_payload(result.nil? ? 204 : 200, result)
       else
-        logger.error("No route found")
-        # This is not something we can deal with, pass it on
-        @code    = 404
-        @payload = ""
+        update_status_and_payload(404)
       end
-      update_content_length
+
+    rescue HTTPException => e # Deal with HTTP exceptions
+      logger.error(e.message)
+      headers.update(e.headers)
+      update_status_and_payload(e.status, e.payload)
     end
 
     def logger
       Restfulness.logger
     end
 
     protected
+
+    def update_status_and_payload(status, payload = "")
+      self.status  = status
+      self.payload = payload
+    end
+
+    def status=(code)
+      @status = code
+    end
+
+    def payload=(body)
+      if body.nil? || body.is_a?(String)
+        @payload = body.to_s
+        update_content_headers(:text)
+      else
+        @payload = MultiJson.encode(body)
+        update_content_headers(:json)
+      end
+    end
     
-    def update_content_length
-      @headers['Content-Length'] = @payload.bytesize.to_s
+    def update_content_headers(type = :json)
+      case type
+      when :json
+        headers['Content-Type'] = 'application/json; charset=utf-8'
+      else # Assume text
+        headers['Content-Type'] = 'text/plain; charset=utf-8'
+      end
+      headers['Content-Length'] = payload.to_s.bytesize.to_s
     end
 
   end
 
 end