2 files changed, 11 insertions, 7 deletions
diff --git a/activerecord/lib/active_record/associations/belongs_to_association.rb b/activerecord/lib/active_record/associations/belongs_to_association.rb
index 5ddff13711..7d96f0c372 100644
--- a/activerecord/lib/active_record/associations/belongs_to_association.rb
+++ b/activerecord/lib/active_record/associations/belongs_to_association.rb
@@ -11,7 +11,11 @@ module ActiveRecord
         raise_on_type_mismatch!(record) if record
 
         update_counters(record)
-        replace_keys(record)
+        if record
+          replace_keys(record)
+        else
+          remove_keys
+        end
         set_inverse_instance(record) if record
 
         @updated = true if record
@@ -58,11 +62,11 @@ module ActiveRecord
         end
 
         def replace_keys(record)
-          if record
-            owner[reflection.foreign_key] = record[reflection.association_primary_key(record.class)]
-          else
-            owner[reflection.foreign_key] = nil
-          end
+          owner[reflection.foreign_key] = record[reflection.association_primary_key(record.class)]
+        end
+
+        def remove_keys
+          owner[reflection.foreign_key] = 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 eae5eed3a1..81d4abfa68 100644
--- a/activerecord/lib/active_record/associations/belongs_to_polymorphic_association.rb
+++ b/activerecord/lib/active_record/associations/belongs_to_polymorphic_association.rb
@@ -11,7 +11,7 @@ module ActiveRecord
 
         def replace_keys(record)
           super
-          owner[reflection.foreign_type] = record && record.class.base_class.name
+          owner[reflection.foreign_type] = record.class.base_class.name
         end
 
         def different_target?(record)