1 files changed, 20 insertions, 9 deletions
diff --git a/activerecord/test/cases/attribute_decorators_test.rb b/activerecord/test/cases/attribute_decorators_test.rb
index 35393753a2..53bd58e22e 100644
--- a/activerecord/test/cases/attribute_decorators_test.rb
+++ b/activerecord/test/cases/attribute_decorators_test.rb
@@ -28,7 +28,7 @@ module ActiveRecord
 
     teardown do
       return unless @connection
-      @connection.execute 'DROP TABLE IF EXISTS attribute_decorators_model'
+      @connection.drop_table 'attribute_decorators_model' if @connection.table_exists? 'attribute_decorators_model'
       Model.attribute_type_decorations.clear
       Model.reset_column_information
     end
@@ -44,8 +44,8 @@ module ActiveRecord
     end
 
     test "decoration does not eagerly load existing columns" do
+      Model.reset_column_information
       assert_no_queries do
-        Model.reset_column_information
         Model.decorate_attribute_type(:a_string, :test) { |t| StringDecorator.new(t) }
       end
     end
@@ -98,17 +98,28 @@ module ActiveRecord
       assert_equal 'Hello! decorated!', model.a_string
       assert_equal 'whatever', model.another_string
       assert_equal 'Hello! decorated! decorated!', child.a_string
-      # We are round tripping the default, and we don't undo our decoration
-      assert_equal 'whatever decorated! decorated!', child.another_string
+      assert_equal 'whatever decorated!', child.another_string
     end
 
-    test "defaults are decorated on the column" do
-      Model.attribute :a_string, Type::String.new, default: 'whatever'
-      Model.decorate_attribute_type(:a_string, :test) { |t| StringDecorator.new(t) }
+    class Multiplier < SimpleDelegator
+      def type_cast_from_user(value)
+        return if value.nil?
+        value * 2
+      end
+      alias type_cast_from_database type_cast_from_user
+    end
+
+    test "decorating with a proc" do
+      Model.attribute :an_int, Type::Integer.new
+      type_is_integer = proc { |_, type| type.type == :integer }
+      Model.decorate_matching_attribute_types type_is_integer, :multiplier do |type|
+        Multiplier.new(type)
+      end
 
-      column = Model.columns_hash['a_string']
+      model = Model.new(a_string: 'whatever', an_int: 1)
 
-      assert_equal 'whatever decorated!', column.default
+      assert_equal 'whatever', model.a_string
+      assert_equal 2, model.an_int
     end
   end
 end