diff --git a/lib/generators/shortener/shortener_generator.rb b/lib/generators/shortener/shortener_generator.rb
index d6faff3..1de9f5a 100644
--- a/lib/generators/shortener/shortener_generator.rb
+++ b/lib/generators/shortener/shortener_generator.rb
@@ -9,7 +9,7 @@ def self.source_root
   end
 
   def self.next_migration_number(dirname)
-    if ActiveRecord::Base.timestamped_migrations
+    if ActiveRecord.respond_to?(:timestamped_migrations) ? ActiveRecord.timestamped_migrations : ActiveRecord::Base.timestamped_migrations
       Time.new.utc.strftime("%Y%m%d%H%M%S")
     else
       "%.3d" % (current_migration_number(dirname) + 1)
diff --git a/spec/generators/shortener_generator_spec.rb b/spec/generators/shortener_generator_spec.rb
new file mode 100644
index 0000000..a16302d
--- /dev/null
+++ b/spec/generators/shortener_generator_spec.rb
@@ -0,0 +1,10 @@
+require "spec_helper"
+require File.expand_path("../../../lib/generators/shortener/shortener_generator.rb", __FILE__)
+
+RSpec.describe ShortenerGenerator do
+  describe ".next_migration_number" do
+    it "returns the next migration number" do
+      expect { described_class.next_migration_number(".") }.not_to raise_error
+    end
+  end
+end