Also allow ActiveRecord::AssociationRelation

Signed-off-by: Jordan Hollinger <[email protected]>

lib/blueprinter-activerecord/preloader.rb

@@ -28,16 +28,17 @@ 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
     # intelligently handles them. There are several unit tests which confirm this behavior.
     #
     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)

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.