lib/chillout/railtie.rb in chillout-0.8.7 vs lib/chillout/railtie.rb in chillout-0.8.8

@@ -2,19 +2,28 @@
 require 'chillout/listener_injector'
 
 module Chillout
   class Railtie < Rails::Railtie
     config.chillout = ActiveSupport::OrderedOptions.new
+    chillout_init = nil
+
     initializer "chillout.creations_listener_initialization" do |rails_app|
       chillout_config = rails_app.config.chillout
       if chillout_config.present?
-        RailsInitializer.new(rails_app, chillout_config, Rails.logger).start
+        chillout_init = RailsInitializer.new(rails_app, chillout_config, Rails.logger)
+        chillout_init.start
       else
         Rails.logger.debug "[Chillout] Not enabled for #{Rails.env}"
       end
     end
 
+    initializer "chillout.after_eager_load", after: :eager_load! do
+      if chillout_init
+        chillout_init.after_eager_load
+      end
+    end
+
     rake_tasks do
       load "chillout/tasks.rb"
     end
 
     generators do
@@ -49,9 +58,14 @@
         @rails_logger.info "[Chillout] Requests monitor enabled"
         Subscribers::ActionControllerNotifications.new.enable(client)
       end
       client.start
       Chillout.client = client
+    end
+
+    def after_eager_load
+      state_machine = Integrations::StateMachine.new
+      state_machine.enable(Chillout.client) if state_machine.available?
     end
 
     private
 
     def options