diff --git a/lib/formtastic/helpers/input_helper.rb b/lib/formtastic/helpers/input_helper.rb
index 4ea9b2c3f..aa9cff0a4 100644
--- a/lib/formtastic/helpers/input_helper.rb
+++ b/lib/formtastic/helpers/input_helper.rb
@@ -299,6 +299,8 @@ def default_input_type(method, options = {}) # @private
       end
 
       # Get a column object for a specified attribute method - if possible.
+      # @return [ActiveModel::Type::Value, #type] in case of rails 5 attributes api
+      # @return [ActiveRecord::ConnectionAdapters::Column] in case of rails 4
       def column_for(method) # @private
         case
           when @object.class.respond_to?(:type_for_attribute)
diff --git a/script/integration-template.rb b/script/integration-template.rb
index 7704f251b..2b3cd988a 100644
--- a/script/integration-template.rb
+++ b/script/integration-template.rb
@@ -18,21 +18,21 @@ def run_bundle
 end
 
 formtastic = -> do
-  generate(:scaffold, 'user name:string password_digest:string')
+  generate(:scaffold, 'user name:string password:digest')
   generate('formtastic:install')
   generate('formtastic:form', 'user name password:password --force')
 
   rake('db:migrate')
 
   in_root do
-    inject_into_file 'app/models/user.rb', "  has_secure_password\n", after: "< ActiveRecord::Base\n"
+    inject_into_class 'app/models/user.rb', 'User', "  has_secure_password\n"
     inject_into_file 'app/assets/stylesheets/application.css', " *= require formtastic\n", before: ' *= require_self'
-    inject_into_file 'test/controllers/users_controller_test.rb', <<-RUBY, before: '  test "should get edit" do'
+    inject_into_class 'test/controllers/users_controller_test.rb', 'UsersControllerTest', <<-RUBY
 
     test "should show form" do
       get :edit, id: @user
 
-      assert_select "form" do
+      assert_select 'form' do
         assert_select 'li.input.string' do
           assert_select 'input#user_name[type=text]'
         end
@@ -46,7 +46,7 @@ def run_bundle
           end
         end
       end
-    end
+    end # test "should show form"
     RUBY
   end
 
diff --git a/script/integration/rails-5-0.rb b/script/integration/rails-5-0.rb
new file mode 100644
index 000000000..6d899a869
--- /dev/null
+++ b/script/integration/rails-5-0.rb
@@ -0,0 +1,22 @@
+gsub_file 'test/controllers/users_controller_test.rb', 'get :edit, id: @user', 'get edit_user_url(@user)'
+
+inject_into_class 'app/models/user.rb', 'User', <<-RUBY
+  class TokenType < ActiveRecord::Type::String
+    def type; :password; end
+  end
+
+  attribute :token, TokenType.new
+RUBY
+
+inject_into_file 'app/views/users/_form.html.erb', <<-HTML, after: '<%= f.input :password %>'
+  <%= f.input :token %>
+HTML
+
+inject_into_file 'test/controllers/users_controller_test.rb', <<-RUBY, before: 'end # test "should show form"'
+    assert_select 'li#user_token_input.password' do
+      assert_select 'input#user_token[type=password]'
+    end
+RUBY
+
+generate 'migration', 'AddTokenToUsers token:string'
+rake 'db:migrate'