lib/multiverse/railtie.rb in multiverse-0.2.0 vs lib/multiverse/railtie.rb in multiverse-0.2.1

@@ -1,26 +1,24 @@
 require "rails/railtie"
 
 module Multiverse
   class Railtie < Rails::Railtie
     generators do
-      if ActiveRecord::VERSION::MAJOR >= 5
-        require "rails/generators/active_record/migration"
-        ActiveRecord::Generators::Migration.prepend(Multiverse::Generators::Migration)
-      else
-        require "rails/generators/migration"
-        Rails::Generators::Migration.prepend(Multiverse::Generators::MigrationTemplate)
-      end
+      ActiveSupport.on_load(:active_record) do
+        if ActiveRecord::VERSION::MAJOR >= 5
+          require "rails/generators/active_record/migration"
+          ActiveRecord::Generators::Migration.prepend(Multiverse::Generators::Migration)
+        else
+          require "rails/generators/migration"
+          Rails::Generators::Migration.prepend(Multiverse::Generators::MigrationTemplate)
+        end
 
-      require "rails/generators/active_record/model/model_generator"
-      ActiveRecord::Generators::ModelGenerator.prepend(Multiverse::Generators::ModelGenerator)
+        require "rails/generators/active_record/model/model_generator"
+        ActiveRecord::Generators::ModelGenerator.prepend(Multiverse::Generators::ModelGenerator)
+        ActiveRecord::Generators::ModelGenerator.prepend(Multiverse::Generators::MigrationTemplate)
 
-      # for Rails < 5.0.3, need to patch db_migrate_path in model and migration generator
-      if ActiveRecord::VERSION::MAJOR == 5 && ActiveRecord.version < Gem::Version.new("5.0.3")
-        ActiveRecord::Generators::ModelGenerator.prepend(Multiverse::Generators::Migration)
-
         require "rails/generators/active_record/migration/migration_generator"
-        ActiveRecord::Generators::MigrationGenerator.prepend(Multiverse::Generators::Migration)
+        ActiveRecord::Generators::MigrationGenerator.prepend(Multiverse::Generators::MigrationTemplate)
       end
     end
 
     rake_tasks do
       namespace :multiverse do