association reloading

lib/has_aux_table.rb

@@ -304,17 +304,31 @@ module HasAuxTable
 
     sig { params(config: AuxTableConfig).void }
     def setup_reload_hook!(config)
+      refresh_object =
+        lambda do |target, fresh|
+          target.instance_variable_set(
+            :@association_cache,
+            fresh.instance_variable_get(:@association_cache)
+          )
+
+          target
+            .instance_variable_get(:@association_cache)
+            .each_value { |association| association.owner = target }
+
+          target.instance_variable_set(
+            :@attributes,
+            fresh.instance_variable_get(:@attributes)
+          )
+        end
+
       self.define_method(:reload) do |*args|
         T.bind(self, ActiveRecord::Base)
         aux_model = config.aux_model_for(self)
-        fresh_model = self.class.find(id)
-        @attributes = fresh_model.instance_variable_get(:@attributes)
-        aux_model.instance_variable_set(
-          :@attributes,
-          fresh_model
-            .association(config.aux_association_name)
-            .target
-            .instance_variable_get(:@attributes)
+        fresh_object = self.class.find(id)
+        refresh_object.call(self, fresh_object)
+        refresh_object.call(
+          aux_model,
+          fresh_object.association(config.aux_association_name).target
         )
         self
       end

spec/active_record/has_aux_table_spec.rb

@@ -758,6 +758,18 @@ RSpec.describe HasAuxTable do
         num_queries = SpecHelper.count_queries { @car.reload }
         expect(num_queries).to eq(1)
       end
+
+      it "reloads associations" do
+        expect(@car.drivers.length).to eq(0)
+
+        Driver.create!(car: @car, name: "Billy Kid")
+        expect(@car.drivers.length).to eq(0)
+        expect(@car.drivers.count).to eq(1)
+
+        @car.reload
+        expect(@car.drivers.length).to eq(1)
+        expect(@car.drivers.count).to eq(1)
+      end
     end
 
     describe "#exists?" do