2 files changed, 8 insertions, 13 deletions
diff --git a/activerecord/lib/active_record/associations/belongs_to_association.rb b/activerecord/lib/active_record/associations/belongs_to_association.rb
index 68c608df13..bd2012df84 100644
--- a/activerecord/lib/active_record/associations/belongs_to_association.rb
+++ b/activerecord/lib/active_record/associations/belongs_to_association.rb
@@ -12,17 +12,20 @@ module ActiveRecord
         if record
           raise_on_type_mismatch!(record)
           update_counters_on_replace(record)
-          replace_keys(record)
           set_inverse_instance(record)
           @updated = true
         else
           decrement_counters
-          remove_keys
         end
 
         self.target = record
       end
 
+      def target=(record)
+        replace_keys(record)
+        super
+      end
+
       def default(&block)
         writer(owner.instance_exec(&block)) if reader.nil?
       end
@@ -78,11 +81,8 @@ module ActiveRecord
         end
 
         def replace_keys(record)
-          owner[reflection.foreign_key] = record._read_attribute(reflection.association_primary_key(record.class))
-        end
-
-        def remove_keys
-          owner[reflection.foreign_key] = nil
+          owner[reflection.foreign_key] = record ?
+            record._read_attribute(reflection.association_primary_key(record.class)) : nil
         end
 
         def foreign_key_present?
diff --git a/activerecord/lib/active_record/associations/belongs_to_polymorphic_association.rb b/activerecord/lib/active_record/associations/belongs_to_polymorphic_association.rb
index 4ce3474bd5..55d789c66a 100644
--- a/activerecord/lib/active_record/associations/belongs_to_polymorphic_association.rb
+++ b/activerecord/lib/active_record/associations/belongs_to_polymorphic_association.rb
@@ -13,12 +13,7 @@ module ActiveRecord
 
         def replace_keys(record)
           super
-          owner[reflection.foreign_type] = record.class.base_class.name
-        end
-
-        def remove_keys
-          super
-          owner[reflection.foreign_type] = nil
+          owner[reflection.foreign_type] = record ? record.class.base_class.name : nil
         end
 
         def different_target?(record)