flash.rb in voltron-flash-0.1.3

- old
+ new

@@ -9,18 +9,26 @@
       base.send :after_action, :include_flash_later
     end
 
     def render(*args)
       include_flash_now
-      super *args
+      super
     end
 
+    def redirect_to(options={}, response_status={})
+      include_flash_later
+      super
+    end
+
     def flash!(**flashes)
       flashes.symbolize_keys.each do |type,messages|
         stored_flashes[type] ||= []
-        stored_flashes[type] << { ajax: flashes.delete(:ajax), messages: Array.wrap(messages) }
+        stored_flashes[type] += Array.wrap(messages)
       end
+
+      # Set the headers initially. If redirecting, they will be removed as the flash will instead be a part of `flash`
+      response.headers[Voltron.config.flash.header] = stored_flashes.to_json
     end
 
     private
 
       def stored_flashes
@@ -28,34 +36,18 @@
       end
 
       # Before rendering, include any flash messages in flash.now,
       # so they will be available when the page is rendered
       def include_flash_now
-        flash_hash(true).each { |type,messages| flash.now[type] = messages }
+        stored_flashes.each { |type,messages| flash.now[type] = messages }
       end
 
-      # If request is an ajax request, or we are redirecting, include flash messages
-      # in the appropriate outlet, either response headers or `flash` itself
+      # When redirecting, remove the flash from the headers (unless ajax request), append it all to the `flash` object
       def include_flash_later
-        if is_redirecting?
-          flash_hash.each { |type,messages| flash[type] = messages }
-        elsif request.xhr?
-          response.headers[Voltron.config.flash.header] = flash_hash.to_json
+        unless request.xhr?
+          response.headers.except! Voltron.config.flash.header
+          stored_flashes.each { |type,messages| flash[type] = messages }
         end
-      end
-
-      def is_redirecting?
-        self.status == 302 || self.status == 301
-      end
-
-      def flash_hash(rendering=false)
-        flashes = stored_flashes.map do |type,messages|
-          { type => messages.map do |f|
-              f[:messages] if !(f[:ajax] == false && request.xhr?) || (f[:ajax] == false && request.xhr? && rendering)
-            end.compact.flatten
-          }
-        end
-        flashes.reduce(Hash.new, :merge).reject { |k,v| v.blank? || k == :ajax }
       end
 
   end
 end
 
\ No newline at end of file