Check that entire collection has been loaded before short circuiting

Currently, when checking that the collection has been loaded, only the first
record is checked. In specific scenarios, if a record is fetched via an `after_initialize`
hook, there is a chance that the first record has been loaded, but other records in the
collection have not.

In order to successfully short circuit the fetching of data, we need to verify that
all of the records in the collection have been loaded.

* Create test for edge case
* Move `reset_callbacks` method to `cases/helper`, since it has been defined in multiple
  locations.

Closes #37730

activerecord/lib/active_record/associations/preloader.rb

@@ -184,7 +184,7 @@ module ActiveRecord
         # and attach it to a relation. The class returned implements a `run` method
         # that accepts a preloader.
         def preloader_for(reflection, owners)
-          if owners.first.association(reflection.name).loaded?
+          if owners.all? { |o| o.association(reflection.name).loaded? }
             return AlreadyLoaded
           end
           reflection.check_preloadable!

activerecord/test/cases/associations/eager_test.rb

@@ -626,6 +626,21 @@ class EagerAssociationTest < ActiveRecord::TestCase
     assert_equal [comments(:does_it_hurt)], assert_no_queries { author.special_post_comments }
   end
 
+  def test_preloading_with_has_one_through_an_sti_with_after_initialize
+    author_a = Author.create!(name: "A")
+    author_b = Author.create!(name: "B")
+    post_a = StiPost.create!(author: author_a, title: "TITLE", body: "BODY")
+    post_b = SpecialPost.create!(author: author_b, title: "TITLE", body: "BODY")
+    comment_a = SpecialComment.create!(post: post_a, body: "TEST")
+    comment_b = SpecialComment.create!(post: post_b, body: "TEST")
+    reset_callbacks(StiPost, :initialize) do
+      StiPost.after_initialize { author }
+      comments = SpecialComment.where(id: [comment_a.id, comment_b.id]).includes(:author).to_a
+      comments_with_author = comments.map { |c| [c.id, c.author.try(:id)] }
+      assert_equal comments_with_author.size, comments_with_author.map(&:second).compact.size
+    end
+  end
+
   def test_preloading_has_many_through_with_implicit_source
     authors = Author.includes(:very_special_comments).to_a
     assert_no_queries do

activerecord/test/cases/associations/has_many_associations_test.rb

@@ -2831,7 +2831,7 @@ class HasManyAssociationsTest < ActiveRecord::TestCase
   end
 
   test "prevent double insertion of new object when the parent association loaded in the after save callback" do
-    reset_callbacks(:save, Bulb) do
+    reset_callbacks(Bulb, :save) do
       Bulb.after_save { |record| record.car.bulbs.load }
 
       car = Car.create!
@@ -2842,7 +2842,7 @@ class HasManyAssociationsTest < ActiveRecord::TestCase
   end
 
   test "prevent double firing the before save callback of new object when the parent association saved in the callback" do
-    reset_callbacks(:save, Bulb) do
+    reset_callbacks(Bulb, :save) do
       count = 0
       Bulb.before_save { |record| record.car.save && count += 1 }
 
@@ -2961,7 +2961,7 @@ class HasManyAssociationsTest < ActiveRecord::TestCase
   end
 
   def test_loading_association_in_validate_callback_doesnt_affect_persistence
-    reset_callbacks(:validation, Bulb) do
+    reset_callbacks(Bulb, :validation) do
       Bulb.after_validation { |record| record.car.bulbs.load }
 
       car = Car.create!(name: "Car")
@@ -2996,18 +2996,4 @@ class HasManyAssociationsTest < ActiveRecord::TestCase
     def force_signal37_to_load_all_clients_of_firm
       companies(:first_firm).clients_of_firm.load_target
     end
-
-    def reset_callbacks(kind, klass)
-      old_callbacks = {}
-      old_callbacks[klass] = klass.send("_#{kind}_callbacks").dup
-      klass.subclasses.each do |subclass|
-        old_callbacks[subclass] = subclass.send("_#{kind}_callbacks").dup
-      end
-      yield
-    ensure
-      klass.send("_#{kind}_callbacks=", old_callbacks[klass])
-      klass.subclasses.each do |subclass|
-        subclass.send("_#{kind}_callbacks=", old_callbacks[subclass])
-      end
-    end
 end

activerecord/test/cases/associations/inverse_associations_test.rb

@@ -545,13 +545,6 @@ class InverseHasManyTests < ActiveRecord::TestCase
       assert_predicate Man.joins(:interests).includes(:interests).first.interests, :any?
     end
   end
-
-  def reset_callbacks(target, type)
-    old_callbacks = target.send(:get_callbacks, type).deep_dup
-    yield
-  ensure
-    target.send(:set_callbacks, type, old_callbacks) if old_callbacks
-  end
 end
 
 class InverseBelongsToTests < ActiveRecord::TestCase

activerecord/test/cases/test_case.rb

@@ -87,6 +87,20 @@ module ActiveRecord
     ensure
       ActiveRecord::Base.has_many_inversing = old
     end
+
+    def reset_callbacks(klass, kind)
+      old_callbacks = {}
+      old_callbacks[klass] = klass.send("_#{kind}_callbacks").dup
+      klass.subclasses.each do |subclass|
+        old_callbacks[subclass] = subclass.send("_#{kind}_callbacks").dup
+      end
+      yield
+    ensure
+      klass.send("_#{kind}_callbacks=", old_callbacks[klass])
+      klass.subclasses.each do |subclass|
+        subclass.send("_#{kind}_callbacks=", old_callbacks[subclass])
+      end
+    end
   end
 
   class PostgreSQLTestCase < TestCase

activerecord/test/models/comment.rb

@@ -59,6 +59,7 @@ class Comment < ActiveRecord::Base
 end
 
 class SpecialComment < Comment
+  has_one :author, through: :post
   default_scope { where(deleted_at: nil) }
 
   def self.what_are_you