diff --git a/closure_tree.gemspec b/closure_tree.gemspec
index 1c7686bb..e950dff3 100644
--- a/closure_tree.gemspec
+++ b/closure_tree.gemspec
@@ -19,13 +19,13 @@ Gem::Specification.new do |gem|
   gem.add_runtime_dependency 'activerecord', '>= 4.1.0'
   gem.add_runtime_dependency 'with_advisory_lock', '>= 3.0.0'
 
+  gem.add_development_dependency 'appraisal'
+  gem.add_development_dependency 'database_cleaner'
+  gem.add_development_dependency 'generator_spec'
+  gem.add_development_dependency 'parallel'
   gem.add_development_dependency 'rspec-instafail'
   gem.add_development_dependency 'rspec-rails'
-  gem.add_development_dependency 'database_cleaner'
-  gem.add_development_dependency 'appraisal'
   gem.add_development_dependency 'timecop'
-  gem.add_development_dependency 'parallel'
-  # gem.add_development_dependency 'ammeter', '1.1.2' # See https://github.com/mceachen/closure_tree/issues/181
   # gem.add_development_dependency 'byebug'
   # gem.add_development_dependency 'ruby-prof' # <- don't need this normally.
 end
diff --git a/lib/generators/closure_tree/migration_generator.rb b/lib/generators/closure_tree/migration_generator.rb
index e444d2f1..fc1f0cf8 100644
--- a/lib/generators/closure_tree/migration_generator.rb
+++ b/lib/generators/closure_tree/migration_generator.rb
@@ -1,5 +1,6 @@
 require 'closure_tree/active_record_support'
 require 'forwardable'
+require 'rails/generators'
 require 'rails/generators/active_record'
 require 'rails/generators/named_base'
 
@@ -41,6 +42,13 @@ def ct
         end
       end
 
+      def migration_version
+        major = ActiveRecord::VERSION::MAJOR
+        if major >= 5
+          "[#{major}.#{ActiveRecord::VERSION::MINOR}]"
+        end
+      end
+
       def self.next_migration_number(dirname)
         ActiveRecord::Generators::Base.next_migration_number(dirname)
       end
diff --git a/lib/generators/closure_tree/templates/create_hierarchies_table.rb.erb b/lib/generators/closure_tree/templates/create_hierarchies_table.rb.erb
index 98af4f4c..18d7671c 100644
--- a/lib/generators/closure_tree/templates/create_hierarchies_table.rb.erb
+++ b/lib/generators/closure_tree/templates/create_hierarchies_table.rb.erb
@@ -1,4 +1,4 @@
-class <%= migration_class_name %> < ActiveRecord::Migration
+class <%= migration_class_name %> < ActiveRecord::Migration<%= migration_version %>
   def change
     create_table :<%= migration_name %>, id: false do |t|
       t.<%= primary_key_type %> :ancestor_id, null: false
diff --git a/spec/generators/migration_generator_spec.rb b/spec/generators/migration_generator_spec.rb
index f3e0e510..f0f20858 100644
--- a/spec/generators/migration_generator_spec.rb
+++ b/spec/generators/migration_generator_spec.rb
@@ -1,48 +1,46 @@
-# require 'spec_helper'
-# require 'ammeter/init'
-#
-# # Generators are not automatically loaded by Rails
-# require 'generators/closure_tree/migration_generator'
-#
-# RSpec.describe ClosureTree::Generators::MigrationGenerator, type: :generator do
-#   TMPDIR = Dir.mktmpdir
-#   # Tell generator where to put its output
-#   destination TMPDIR
-#   before { prepare_destination }
-#
-#   describe 'generator output' do
-#     before { run_generator %w(tag) }
-#     subject { migration_file('db/migrate/create_tag_hierarchies.rb') }
-#     it { is_expected.to be_a_migration }
-#     it { is_expected.to contain(/t.integer :ancestor_id, null: false/) }
-#     it { is_expected.to contain(/t.integer :descendant_id, null: false/) }
-#     it { is_expected.to contain(/t.integer :generations, null: false/) }
-#     it { is_expected.to contain(/add_index :tag_hierarchies/) }
-#   end
-#
-#   describe 'generator output with namespaced model' do
-#     before { run_generator %w(Namespace::Type) }
-#     subject { migration_file('db/migrate/create_namespace_type_hierarchies.rb') }
-#     it { is_expected.to be_a_migration }
-#     it { is_expected.to contain(/t.integer :ancestor_id, null: false/) }
-#     it { is_expected.to contain(/t.integer :descendant_id, null: false/) }
-#     it { is_expected.to contain(/t.integer :generations, null: false/) }
-#     it { is_expected.to contain(/add_index :namespace_type_hierarchies/) }
-#   end
-#
-#   describe 'generator output with namespaced model with /' do
-#     before { run_generator %w(namespace/type) }
-#     subject { migration_file('db/migrate/create_namespace_type_hierarchies.rb') }
-#     it { is_expected.to be_a_migration }
-#     it { is_expected.to contain(/t.integer :ancestor_id, null: false/) }
-#     it { is_expected.to contain(/t.integer :descendant_id, null: false/) }
-#     it { is_expected.to contain(/t.integer :generations, null: false/) }
-#     it { is_expected.to contain(/add_index :namespace_type_hierarchies/) }
-#   end
-#
-#   it 'should run all tasks in generator without errors' do
-#     gen = generator %w(tag)
-#     expect(gen).to receive :create_migration_file
-#     capture(:stdout) { gen.invoke_all }
-#   end
-# end
+require 'spec_helper'
+require "generator_spec/test_case"
+
+# Generators are not automatically loaded by Rails
+require 'generators/closure_tree/migration_generator'
+
+RSpec.describe ClosureTree::Generators::MigrationGenerator, type: :generator do
+  include GeneratorSpec::TestCase
+
+  # Tell generator where to put its output
+  destination Dir.mktmpdir
+  before { prepare_destination }
+
+  describe 'generator output' do
+    before { run_generator %w(tag) }
+    subject { File.read migration_file_name('db/migrate/create_tag_hierarchies.rb') }
+    it { is_expected.to match(/t.integer :ancestor_id, null: false/) }
+    it { is_expected.to match(/t.integer :descendant_id, null: false/) }
+    it { is_expected.to match(/t.integer :generations, null: false/) }
+    it { is_expected.to match(/add_index :tag_hierarchies/) }
+  end
+
+  describe 'generator output with namespaced model' do
+    before { run_generator %w(Namespace::Type) }
+    subject { File.read migration_file_name('db/migrate/create_namespace_type_hierarchies.rb') }
+    it { is_expected.to match(/t.integer :ancestor_id, null: false/) }
+    it { is_expected.to match(/t.integer :descendant_id, null: false/) }
+    it { is_expected.to match(/t.integer :generations, null: false/) }
+    it { is_expected.to match(/add_index :namespace_type_hierarchies/) }
+  end
+
+  describe 'generator output with namespaced model with /' do
+    before { run_generator %w(namespace/type) }
+    subject { File.read migration_file_name('db/migrate/create_namespace_type_hierarchies.rb') }
+    it { is_expected.to match(/t.integer :ancestor_id, null: false/) }
+    it { is_expected.to match(/t.integer :descendant_id, null: false/) }
+    it { is_expected.to match(/t.integer :generations, null: false/) }
+    it { is_expected.to match(/add_index :namespace_type_hierarchies/) }
+  end
+
+  it 'should run all tasks in generator without errors' do
+    gen = generator %w(tag)
+    expect(gen).to receive :create_migration_file
+    capture(:stdout) { gen.invoke_all }
+  end
+end