Add spec for ar intergration and context (#62)

* Explicitly show no context set when using AR preloading
* Raise error when n1_bind_to received invalid arguments
* Refactor syntax and fix possible thread safe issue
* Fix context propagation when using n1_bind_to

Author: djezzzl

CHANGELOG.md

@@ -1,3 +1,9 @@
+## [Unreleased]
+
+- Fix a rare thread-safety issue for setting context 
+- Raise errors if `n1_bind_to` received unexpected arguments
+- Fix context propagation to loaded objects when using `n1_bind_to`
+
 ## [2.2.0] - 2026/03/16
 
 - Support both identity and equality comparison in `Loader#for`. Identity lookup (via `object_id`) is tried first, with equality lookup as a fallback. Thanks [Alfonso Uceda](https://github.com/AlfonsoUceda) for reporting the issue!

lib/n1_loader.rb

@@ -15,4 +15,5 @@ class NotLoaded < Error; end
   class NotFilled < Error; end
   class MissingArgument < Error; end
   class InvalidArgument < Error; end
+  class InvalidBinding < Error; end
 end

lib/n1_loader/active_record/loader.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 N1Loader::Loader.define_method :preloaded_records do
-  @preloaded_records ||= loaded.values.flatten
+  @preloaded_records ||= loaded? && loaded_by_value.values.flatten
 end

lib/n1_loader/ar_lazy_preload/loader_patch.rb

@@ -6,17 +6,22 @@ module ArLazyPreload
     module LoaderPatch
       attr_accessor :context_setup
 
-      def loaded
-        return @loaded_by_identity if @already_loaded && @already_context
+      def loaded?
+        return true if @already_loaded && @already_context
 
         super
 
-        synchronize do
-          context_setup&.call(@loaded_by_identity.values.flatten) unless @already_context
-        end
+        synchronize { non_thread_safe_context_setting unless @already_context }
+
+        true
+      end
+
+      def non_thread_safe_context_setting
+        return if @already_context
+
+        context_setup&.call(loaded_by_identity.values.flatten)
 
         @already_context = true
-        @loaded_by_identity
       end
     end
   end

lib/n1_loader/core/loadable.rb

@@ -33,9 +33,21 @@ def n1_loader(name)
     end
 
     def n1_bind_to(collection)
+      unless collection.is_a?(Array) && collection.any? do |obj|
+        obj == self || obj.equal?(self)
+      end
+
+        raise InvalidBinding,
+              "assigned collection should be array and include object"
+      end
+
       @n1_binding = collection
     end
 
+    def n1_bind_to?
+      !@n1_binding.nil?
+    end
+
     def n1_loader_reload(name)
       elements = @n1_binding || [self]
       collection = LoaderCollection.new(self.class.n1_loaders[name], elements)

lib/n1_loader/core/loader.rb

@@ -47,14 +47,14 @@ def initialize(elements, **args)
     end
 
     def for(element)
-      identity_loaded = loaded
+      return unless loaded?
 
-      if identity_loaded.empty? && elements.any?
+      if loaded_by_identity.empty? && elements.any?
         raise NotFilled, "Nothing was preloaded, perhaps you forgot to use fulfill method"
       end
 
-      return identity_loaded[element] if identity_loaded.key?(element)
-      return @loaded_by_value[element] if @loaded_by_value.key?(element)
+      return loaded_by_identity[element] if loaded_by_identity.key?(element)
+      return loaded_by_value[element] if loaded_by_value.key?(element)
 
       raise NotLoaded, "The data was not preloaded for the given element"
     end
@@ -66,7 +66,7 @@ def cache_key
 
     private
 
-    attr_reader :elements, :args
+    attr_reader :elements, :args, :loaded_by_value, :loaded_by_identity
 
     def check_missing_arguments!
       return unless (arguments = self.class.arguments)
@@ -107,22 +107,19 @@ def perform(_elements)
     end
 
     def fulfill(element, value)
-      @loaded_by_identity[element] = value
-      @loaded_by_value[element] = value
+      loaded_by_identity[element] = value
+      loaded_by_value[element] = value
     end
 
-    def ensure_loaded
-      return if @already_loaded
+    def loaded?
+      return true if @already_loaded
 
-      synchronize { non_thread_safe_loaded unless @already_loaded }
-    end
+      synchronize { non_thread_safe_loading unless @already_loaded }
 
-    def loaded
-      ensure_loaded
-      @loaded_by_identity
+      true
     end
 
-    def non_thread_safe_loaded # rubocop:disable Metrics/AbcSize, Metrics/MethodLength
+    def non_thread_safe_loading # rubocop:disable Metrics/AbcSize, Metrics/MethodLength, Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity
       return if @already_loaded
 
       check_arguments!
@@ -133,8 +130,13 @@ def non_thread_safe_loaded # rubocop:disable Metrics/AbcSize, Metrics/MethodLeng
       if respond_to?(:single) && elements.size == 1
         fulfill(elements.first, single(elements.first))
       elsif elements.any?
-        elements.each { |el| el.n1_bind_to(elements) if el.respond_to?(:n1_bind_to) }
         perform(elements)
+
+        # propagate context to loaded objects only when it was set
+        if elements.first.respond_to?(:n1_bind_to?) && elements.first.n1_bind_to?
+          loaded_objects = loaded_by_identity.values.flatten
+          loaded_objects.each { |el| el.n1_bind_to(loaded_objects) if el.respond_to?(:n1_bind_to) }
+        end
       end
 
       @already_loaded = true

spec/activerecord_spec.rb

@@ -233,6 +233,21 @@ def perform(elements)
         .and change(Entity, :count).by(1)
     end
 
+    context "when preloading AR" do
+      let(:objects) { Entity.includes(:company) }
+
+      it "doesn't set context further for N1Loader" do
+        expect do
+          objects.each do |object|
+            expect(object.company.data).to eq(object)
+          end
+            .to make_database_queries(matching: /companies/, count: 1)
+            .and make_database_queries(matching: /entities/, count: 2)
+            .and make_database_queries(count: 3)
+        end
+      end
+    end
+
     context "with arguments" do
       before { skip "unsupported by ArLazyPreload" if ar_lazy_preload_defined? }
 

spec/n1_loader_spec.rb

@@ -121,6 +121,16 @@ def perform(elements)
           end
         end
       end
+
+      n1_optimized :new_objects do |elements|
+        elements.first.class.perform!
+
+        objects = elements.map { elements.first.class.new }
+
+        elements.each_with_index do |element, index|
+          fulfill(element, objects[index])
+        end
+      end
     end
   end
 
@@ -475,11 +485,14 @@ def perform(elements)
   end
 
   describe "n1_bind_to" do
-    it "returns correct data for each bound object" do
-      objects.each { |obj| obj.n1_bind_to(objects) }
+    it "raises error if invalid collection was passed" do
+      expect do
+        objects.each { |obj| obj.n1_bind_to([]) }
+      end.to raise_error N1Loader::InvalidBinding
 
-      expect(objects.first.inline).to eq([objects.first])
-      expect(objects.last.inline).to eq([objects.last])
+      expect do
+        objects.first.n1_bind_to([objects.last])
+      end.to raise_error N1Loader::InvalidBinding
     end
 
     it "loads all bound objects in a single batch" do
@@ -501,67 +514,31 @@ def perform(elements)
       expect { objects.first.inline }.to change(klass, :count).by(1)
       expect { objects.last.inline }.not_to change(klass, :count)
     end
-  end
-
-  describe "automatic context binding" do
-    let(:nested_klass) do
-      Class.new do
-        include N1Loader::Loadable
-
-        class << self
-          def perform!(loader)
-            @counts ||= {}
-            @counts[loader] = (@counts[loader] || 0) + 1
-          end
 
-          def count(loader)
-            @counts&.fetch(loader, 0) || 0
-          end
-        end
-
-        n1_optimized :first_loader do |elements|
-          elements.first.class.perform!(:first_loader)
-          elements.each { |el| fulfill(el, el) }
-        end
-
-        n1_optimized :second_loader do |elements|
-          elements.first.class.perform!(:second_loader)
-          elements.each { |el| fulfill(el, el) }
-        end
-      end
-    end
-
-    let(:nested_objects) { [nested_klass.new, nested_klass.new] }
-
-    it "automatically sets shared context when loaded through N1Loader" do
-      N1Loader::Preloader.new(nested_objects).preload(:first_loader)
-
-      # Accessing first_loader triggers perform([obj1, obj2]),
-      # which automatically calls n1_bind_to on both objects
-      nested_objects.first.first_loader
-
-      # Both objects are now auto-bound; second_loader should batch in one perform call
-      expect { nested_objects.map(&:second_loader) }.to change { nested_klass.count(:second_loader) }.by(1)
-    end
-
-    it "auto-binds so second access on sibling does not trigger another load" do
-      N1Loader::Preloader.new(nested_objects).preload(:first_loader)
+    it "propagates context to loaded objects" do
+      objects.each { |obj| obj.n1_bind_to(objects) }
 
-      nested_objects.first.first_loader
+      expect do
+        objects.each(&:new_objects)
+      end.to change(klass, :count).by(1)
 
-      nested_objects.first.second_loader
-      expect { nested_objects.last.second_loader }.not_to change { nested_klass.count(:second_loader) } # rubocop:disable Lint/AmbiguousBlockAssociation
+      expect do
+        objects.map(&:new_objects).map(&:inline)
+      end.to change(klass, :count).by(1)
     end
+  end
 
-    it "auto-binds single element to a collection containing only itself" do
-      single = nested_klass.new
-      N1Loader::Preloader.new([single]).preload(:first_loader)
+  describe "with preloader" do
+    it "doesn't propagate context to loaded objects" do
+      N1Loader::Preloader.new(objects).preload(:new_objects)
 
-      single.first_loader
+      expect do
+        objects.each(&:new_objects)
+      end.to change(klass, :count).by(1)
 
-      # Accessing second_loader on the single element should still work correctly
-      expect { single.second_loader }.to change { nested_klass.count(:second_loader) }.by(1)
-      expect { single.second_loader }.not_to change { nested_klass.count(:second_loader) } # rubocop:disable Lint/AmbiguousBlockAssociation
+      expect do
+        objects.map(&:new_objects).map(&:inline)
+      end.to change(klass, :count).by(2)
     end
   end
 end