Make counter cache work when replacing an association (closes #3245). Thanks for the patch!

git-svn-id: http://svn-commit.rubyonrails.org/rails/trunk@3762 5ecf4fe2-1ee6-0310-87b1-e25e094e27de

4 files changed, 53 insertions, 0 deletions

diff --git a/activerecord/CHANGELOG b/activerecord/CHANGELOG
index bcd5f0e430..0dccedcf1f 100644
--- a/activerecord/CHANGELOG
+++ b/activerecord/CHANGELOG
@@ -1,5 +1,7 @@
 *SVN*
 
+* Make counter cache work when replacing an association #3245 [eugenol@gmail.com]
+
 * Make migrations verbose [Jamis Buck]
 
 * Make counter_cache work with polymorphic belongs_to [Jamis Buck]
diff --git a/activerecord/lib/active_record/associations/belongs_to_association.rb b/activerecord/lib/active_record/associations/belongs_to_association.rb
index 804a7ebf21..3e0e0d0fb9 100644
--- a/activerecord/lib/active_record/associations/belongs_to_association.rb
+++ b/activerecord/lib/active_record/associations/belongs_to_association.rb
@@ -10,11 +10,22 @@ module ActiveRecord
       end
 
       def replace(record)
+        counter_cache_name = @reflection.counter_cache_column
+
         if record.nil?
+          if counter_cache_name && @owner[counter_cache_name] && !@owner.new_record?
+            @reflection.klass.decrement_counter(counter_cache_name, @owner[@reflection.primary_key_name]) if @owner[@reflection.primary_key_name]
+          end
+
           @target = @owner[@reflection.primary_key_name] = nil
         else
           raise_on_type_mismatch(record)
 
+          if counter_cache_name && @owner[counter_cache_name] && !@owner.new_record?
+            @reflection.klass.increment_counter(counter_cache_name, record.id)
+            @reflection.klass.decrement_counter(counter_cache_name, @owner[@reflection.primary_key_name]) if @owner[@reflection.primary_key_name]
+          end
+
           @target = (AssociationProxy === record ? record.target : record)
           @owner[@reflection.primary_key_name] = record.id unless record.new_record?
           @updated = true
diff --git a/activerecord/lib/active_record/reflection.rb b/activerecord/lib/active_record/reflection.rb
index b9369feeb7..32e812a86a 100644
--- a/activerecord/lib/active_record/reflection.rb
+++ b/activerecord/lib/active_record/reflection.rb
@@ -131,6 +131,14 @@ module ActiveRecord
         @association_foreign_key ||= @options[:association_foreign_key] || class_name.foreign_key
       end
 
+      def counter_cache_column
+        if options[:counter_cache] == true
+          "#{active_record.name.underscore.pluralize}_count"
+        elsif options[:counter_cache]
+          options[:counter_cache]
+        end
+      end
+
       private
         def name_to_class_name(name)
           if name =~ /::/
diff --git a/activerecord/test/associations_test.rb b/activerecord/test/associations_test.rb
index c6cf7807fa..f687d3781e 100755
--- a/activerecord/test/associations_test.rb
+++ b/activerecord/test/associations_test.rb
@@ -829,6 +829,38 @@ class BelongsToAssociationsTest < Test::Unit::TestCase
     assert_equal 0, Topic.find(debate.id).send(:read_attribute, "replies_count"), "First reply deleted"
   end
 
+  def test_belongs_to_counter_with_reassigning
+    t1 = Topic.create("title" => "t1")
+    t2 = Topic.create("title" => "t2")
+    r1 = Reply.new("title" => "r1", "content" => "r1")
+    r1.topic = t1
+
+    assert r1.save
+    assert_equal 1, Topic.find(t1.id).replies.size
+    assert_equal 0, Topic.find(t2.id).replies.size
+
+    r1.topic = Topic.find(t2.id)
+
+    assert r1.save
+    assert_equal 0, Topic.find(t1.id).replies.size
+    assert_equal 1, Topic.find(t2.id).replies.size
+
+    r1.topic = nil
+
+    assert_equal 0, Topic.find(t1.id).replies.size
+    assert_equal 0, Topic.find(t2.id).replies.size
+
+    r1.topic = t1
+
+    assert_equal 1, Topic.find(t1.id).replies.size
+    assert_equal 0, Topic.find(t2.id).replies.size
+
+    r1.destroy
+
+    assert_equal 0, Topic.find(t1.id).replies.size
+    assert_equal 0, Topic.find(t2.id).replies.size
+  end
+
   def test_assignment_before_parent_saved
     client = Client.find(:first)
     apple = Firm.new("name" => "Apple")