Fix `different_target?` to respect custom primary key counter

2 files changed, 13 insertions, 3 deletions

diff --git a/activerecord/lib/active_record/associations/belongs_to_association.rb b/activerecord/lib/active_record/associations/belongs_to_association.rb
index 775559e224..121592733c 100644
--- a/activerecord/lib/active_record/associations/belongs_to_association.rb
+++ b/activerecord/lib/active_record/associations/belongs_to_association.rb
@@ -86,12 +86,15 @@ module ActiveRecord
 
         # Checks whether record is different to the current target, without loading it
         def different_target?(record)
-          record.id != owner._read_attribute(reflection.foreign_key)
+          record._read_attribute(primary_key(record)) != owner._read_attribute(reflection.foreign_key)
         end
 
         def replace_keys(record)
-          owner[reflection.foreign_key] = record ?
-            record._read_attribute(reflection.association_primary_key(record.class)) : nil
+          owner[reflection.foreign_key] = record ? record._read_attribute(primary_key(record)) : nil
+        end
+
+        def primary_key(record)
+          reflection.association_primary_key(record.class)
         end
 
         def foreign_key_present?
diff --git a/activerecord/test/cases/associations/belongs_to_associations_test.rb b/activerecord/test/cases/associations/belongs_to_associations_test.rb
index b3f5de4c5f..c0dfa91a21 100644
--- a/activerecord/test/cases/associations/belongs_to_associations_test.rb
+++ b/activerecord/test/cases/associations/belongs_to_associations_test.rb
@@ -471,6 +471,13 @@ class BelongsToAssociationsTest < ActiveRecord::TestCase
     assert_equal 1, debate.reload.replies_count
     assert_equal 0, debate2.reload.replies_count
 
+    assert_no_queries do
+      reply.topic_with_primary_key = debate
+    end
+
+    assert_equal 1, debate.reload.replies_count
+    assert_equal 0, debate2.reload.replies_count
+
     reply.topic_with_primary_key = debate2
 
     assert_equal 0, debate.reload.replies_count