From 42d3ce97d93743dd5256550cbdf004d59a501008 Mon Sep 17 00:00:00 2001
From: Dylan Knutson <tcdknutson@gmail.com>
Date: Sat, 19 Jul 2025 00:14:56 +0000
Subject: [PATCH] association reloading

---
 lib/has_aux_table.rb                     | 30 +++++++++++++++++-------
 spec/active_record/has_aux_table_spec.rb | 12 ++++++++++
 2 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/lib/has_aux_table.rb b/lib/has_aux_table.rb
index 920695d..8fdcc68 100644
--- a/lib/has_aux_table.rb
+++ b/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
diff --git a/spec/active_record/has_aux_table_spec.rb b/spec/active_record/has_aux_table_spec.rb
index a7ae1bf..1a602d5 100644
--- a/spec/active_record/has_aux_table_spec.rb
+++ b/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