From e3cbe4f49fe7d5f044f3ecbc0947502b1ea40542 Mon Sep 17 00:00:00 2001
From: Jamis Buck <jamis.buck@mongodb.com>
Date: Tue, 18 Feb 2025 08:36:02 -0700
Subject: [PATCH] MONGOID-5844 Fix counting bug in HABTM associations (#5951)

In certain situations, calling #size on a HABTM association will
return the wrong count.

This will happen if the association is
initialized to a single element (forcing the _unloaded Criteria
object to be scoped to that specific element) and then assigning
an array of multiple (already-persisted) elements to the association,
where one of the elements is the same element that already existed
there.

In this case, `_unloaded.count` will return 1, and then the
_added array will have two previously-persisted records. Naive
implementations will thus return either 1, or 3, rather than the
correct answer of 2. To get the correct answer, it is necessary
to add a filter condition to `_unloaded.count` that excludes the
records in the `_added` array.
---
 .../referenced/has_many/enumerable.rb         | 25 +++++--
 .../has_and_belongs_to_many/proxy_spec.rb     | 65 ++++++++-----------
 2 files changed, 49 insertions(+), 41 deletions(-)

diff --git a/lib/mongoid/association/referenced/has_many/enumerable.rb b/lib/mongoid/association/referenced/has_many/enumerable.rb
index b388e404be..1b88222935 100644
--- a/lib/mongoid/association/referenced/has_many/enumerable.rb
+++ b/lib/mongoid/association/referenced/has_many/enumerable.rb
@@ -419,11 +419,28 @@ def respond_to?(name, include_private = false)
           #
           # @return [ Integer ] The size of the enumerable.
           def size
-            count = (_unloaded ? _unloaded.count : _loaded.count)
-            if count.zero?
-              count + _added.count
+            # If _unloaded is present, then it will match the set of documents
+            # that belong to this association, which have already been persisted
+            # to the database. This set of documents must be considered when
+            # computing the size of the association, along with anything that has
+            # since been added.
+            if _unloaded
+              if _added.any?
+                # Note that _added may include records that _unloaded already
+                # matches. This is the case if the association is assigned an array
+                # of items and some of them were already elements of the association.
+                #
+                # we need to thus make sure _unloaded.count excludes any elements
+                # that already exist in _added.
+
+                count = _unloaded.not(:_id.in => _added.values.map(&:id)).count
+                count + _added.values.count
+              else
+                _unloaded.count
+              end
+
             else
-              count + _added.values.count { |d| d.new_record? }
+              _loaded.count + _added.count
             end
           end
 
diff --git a/spec/mongoid/association/referenced/has_and_belongs_to_many/proxy_spec.rb b/spec/mongoid/association/referenced/has_and_belongs_to_many/proxy_spec.rb
index d8cad40347..18b2691943 100644
--- a/spec/mongoid/association/referenced/has_and_belongs_to_many/proxy_spec.rb
+++ b/spec/mongoid/association/referenced/has_and_belongs_to_many/proxy_spec.rb
@@ -1756,43 +1756,6 @@
       end
     end
 
-    describe "#any?" do
-
-      let(:person) do
-        Person.create!
-      end
-
-      context "when nothing exists on the relation" do
-
-        context "when no document is added" do
-
-          let!(:sandwich) do
-            Sandwich.create!
-          end
-
-          it "returns false" do
-            expect(sandwich.meats.any?).to be false
-          end
-        end
-
-        context "when the document is destroyed" do
-
-          before do
-            Meat.create!
-          end
-
-          let!(:sandwich) do
-            Sandwich.create!
-          end
-
-          it "returns false" do
-            sandwich.destroy
-            expect(sandwich.meats.any?).to be false
-          end
-        end
-      end
-    end
-
     context "when documents have been persisted" do
 
       let!(:preference) do
@@ -3041,6 +3004,34 @@
         end
       end
 
+      # MONGOID-5844
+      #
+      # Specifically, this tests the case where the association is
+      # initialized with a single element (so that Proxy#push does not take
+      # the `concat` route), which causes `reset_unloaded` to be called, which
+      # sets the `_unloaded` Criteria object to match only the specific element
+      # that was given.
+      #
+      # The issue now is that when the events list is updated to be both events,
+      # _unloaded matches one of them already, and the other has previously been
+      # persisted so `new_record?` won't match it. We need to make sure the
+      # `#size` logic properly accounts for this case.
+      context 'when documents have been previously persisted' do
+        let(:person1) { Person.create! }
+        let(:person2) { Person.create! }
+        let(:event1) { Event.create!(administrators: [ person1 ]) }
+        let(:event2) { Event.create!(administrators: [ person2 ]) }
+
+        before do
+          person1.administrated_events = [ event1, event2 ]
+        end
+
+        it 'returns the number of associated documents [MONGOID-5844]' do
+          expect(person1.administrated_events.to_a.size).to eq(2)
+          expect(person1.administrated_events.size).to eq(2)
+        end
+      end
+
       context "when documents have not been persisted" do
 
         before do