lib/rails_warden.rb in rails_warden-0.5.8 vs lib/rails_warden.rb in rails_warden-0.6.0

@@ -1,10 +1,9 @@
 # encoding: utf-8
 require 'warden'
 require 'active_support'
 
-
 $:.unshift File.expand_path(File.dirname(__FILE__))
 require "rails_warden/manager"
 require "rails_warden/rails_settings"
 require "rails_warden/controller_mixin"
 
@@ -87,16 +86,32 @@
     end
   end
 else
   class RailsWarden::Railtie < Rails::Railtie
     include_block = Proc.new {
-      ::ActionController::Base.class_eval do
-        include RailsWarden::Mixins::HelperMethods
-        include RailsWarden::Mixins::ControllerOnlyMethods
+
+      ActiveSupport.on_load(:action_controller) do
+        ::ActionController::Base.class_eval do
+          include RailsWarden::Mixins::HelperMethods
+          include RailsWarden::Mixins::ControllerOnlyMethods
+        end
+        if defined? ::ActionController::API
+          ::ActionController::API.class_eval do
+            include RailsWarden::Mixins::HelperMethods
+
+            if defined? helper
+              helper RailsWarden::Mixins::ControllerOnlyMethods
+            else
+              include RailsWarden::Mixins::ControllerOnlyMethods
+            end
+          end
+        end
       end
 
-      ::ActionView::Base.class_eval do
-        include RailsWarden::Mixins::HelperMethods
+      ActiveSupport.on_load(:action_view) do
+        ::ActionView::Base.class_eval do
+          include RailsWarden::Mixins::HelperMethods
+        end
       end
     }
 
     if respond_to?(:initializer)
       initializer :warden, &include_block