From 1cab8ccfdbc0534edfc04606fe93782e906cbb77 Mon Sep 17 00:00:00 2001
From: Jordan Hollinger <jordan.hollinger@procore.com>
Date: Tue, 28 May 2024 19:29:54 -0400
Subject: [PATCH] Also allow ActiveRecord::AssociationRelation

Signed-off-by: Jordan Hollinger <jordan.hollinger@procore.com>
---
 lib/blueprinter-activerecord/preloader.rb |  7 ++++---
 test/preloader_extension_test.rb          | 20 +++++++++++++++-----
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/lib/blueprinter-activerecord/preloader.rb b/lib/blueprinter-activerecord/preloader.rb
index 39db3d8..59832ba 100644
--- a/lib/blueprinter-activerecord/preloader.rb
+++ b/lib/blueprinter-activerecord/preloader.rb
@@ -28,8 +28,9 @@ def initialize(auto: false, use: :preload, &auto_proc)
 
     #
     # Implements the "pre_render" Blueprinter Extension to preload associations from a view.
-    # If auto is true, all ActiveRecord::Relation objects will be preloaded. If auto is false,
-    # only queries that have called `.preload_blueprint` will be preloaded.
+    # If auto is true, all ActiveRecord::Relation and ActiveRecord::AssociationRelation objects
+    # will be preloaded. If auto is false, only queries that have called `.preload_blueprint`
+    # will be preloaded.
     #
     # NOTE: If auto is on, *don't* be concerned that you'll end up with duplicate preloads. Even if
     # the query ends up with overlapping members in 'preload' and 'includes', ActiveRecord
@@ -37,7 +38,7 @@ def initialize(auto: false, use: :preload, &auto_proc)
     #
     def pre_render(object, blueprint, view, options)
       case object.class.name
-      when "ActiveRecord::Relation"
+      when "ActiveRecord::Relation", "ActiveRecord::AssociationRelation"
         if object.preload_blueprint_method || auto || auto_proc&.call(object, blueprint, view, options) == true
           object.before_preload_blueprint = extract_preloads object
           blueprint_preloads = self.class.preloads(blueprint, view, object.model)
diff --git a/test/preloader_extension_test.rb b/test/preloader_extension_test.rb
index 21a34df..6dbccc7 100644
--- a/test/preloader_extension_test.rb
+++ b/test/preloader_extension_test.rb
@@ -5,15 +5,15 @@
 class PreloaderExtensionTest < Minitest::Test
   def setup
     DatabaseCleaner.start
-    customer = Customer.create!(name: "ACME")
-    project = Project.create!(customer_id: customer.id, name: "Project A")
+    @customer = Customer.create!(name: "ACME")
+    project = Project.create!(customer_id: @customer.id, name: "Project A")
     category = Category.create!(name: "Foo")
     ref_plan = RefurbPlan.create!(name: "Plan A")
     battery1 = LiIonBattery.create!(num_ions: 100, num_other: 100, refurb_plan_id: ref_plan.id)
     battery2 = LeadAcidBattery.create!(num_lead: 100, num_acid: 100)
-    Widget.create!(customer_id: customer.id, project_id: project.id, category_id: category.id, name: "Widget A", battery1: battery1, battery2: battery2)
-    Widget.create!(customer_id: customer.id, project_id: project.id, category_id: category.id, name: "Widget B", battery1: battery1)
-    Widget.create!(customer_id: customer.id, project_id: project.id, category_id: category.id, name: "Widget C", battery1: battery1)
+    Widget.create!(customer_id: @customer.id, project_id: project.id, category_id: category.id, name: "Widget A", battery1: battery1, battery2: battery2)
+    Widget.create!(customer_id: @customer.id, project_id: project.id, category_id: category.id, name: "Widget B", battery1: battery1)
+    Widget.create!(customer_id: @customer.id, project_id: project.id, category_id: category.id, name: "Widget C", battery1: battery1)
     Blueprinter.configure do |config|
       config.extensions << BlueprinterActiveRecord::Preloader.new
     end
@@ -146,6 +146,16 @@ def test_auto_preload_with_existing_preloads
     assert_equal({:battery1=>{:refurb_plan=>{}}, :category=>{}, :project=>{:customer=>{}}}, BlueprinterActiveRecord::Helpers.merge_values(q.values[:preload]))
   end
 
+  def test_auto_preload_with_association_relation
+    ext = BlueprinterActiveRecord::Preloader.new(auto: true)
+    q = @customer.widgets.order(:name).strict_loading
+    q = ext.pre_render(q, WidgetBlueprint, :extended, {})
+
+    assert ext.auto
+    assert_equal :preload, ext.use
+    assert_equal [{:battery1=>{:fake_assoc=>{}, :refurb_plan=>{}}, :battery2=>{:fake_assoc=>{}, :refurb_plan=>{}}, :category=>{}, :project=>{:customer=>{}}}], q.values[:preload]
+  end
+
   def test_auto_preload_with_block_true
     ext = BlueprinterActiveRecord::Preloader.new { |object| true }
     q = Widget.