From ad8ae7945b9e81d5457785195269f83118a78397 Mon Sep 17 00:00:00 2001
From: Dylan Knutson <tcdknutson@gmail.com>
Date: Thu, 24 Jul 2025 05:19:46 +0000
Subject: [PATCH] tests for has_one & belongs_to custom foreign keys

---
 lib/has_aux_table/aux_table_config.rb |  25 ++++--
 spec/has_aux_table_spec.rb            | 111 ++++++++++++++++++++++++++
 spec/spec_models.rb                   |  61 ++++++++++++++
 3 files changed, 190 insertions(+), 7 deletions(-)

diff --git a/lib/has_aux_table/aux_table_config.rb b/lib/has_aux_table/aux_table_config.rb
index 37d492d..1550e7b 100644
--- a/lib/has_aux_table/aux_table_config.rb
+++ b/lib/has_aux_table/aux_table_config.rb
@@ -141,14 +141,25 @@ module HasAuxTable
           # attribute is a column on the aux table
           aux[k] = v
         elsif assoc = self.main.klass.reflect_on_association(k.to_s)
-          # attribute is an association on the main class
-          fk = assoc.association_foreign_key
-          if self.aux.column_names.include?(fk)
-            # the association is a column on the aux table, `v` is
-            # a model, get the primary key of the model
-            aux[fk] = v && v.send(assoc.association_primary_key)
+          if assoc.is_a?(ActiveRecord::Reflection::BelongsToReflection)
+            # attribute is an association on the main class
+            fk = assoc.foreign_key
+            if self.aux.column_names.include?(fk)
+              # the association is a column on the aux table, `v` is
+              # a model, get the primary key of the model
+              aux[fk] = v && v.send(assoc.association_primary_key)
+            else
+              # association is on the main table, `v` is a model,
+              main[k] = v
+            end
+          elsif assoc.is_a?(ActiveRecord::Reflection::HasOneReflection)
+            pk = assoc.active_record_primary_key
+            if self.aux.column_names.include?(pk)
+              aux[pk] = v && v.send(assoc.foreign_key)
+            else
+              main[k] = v
+            end
           else
-            # association is on the main table, `v` is a model,
             main[k] = v
           end
         else
diff --git a/spec/has_aux_table_spec.rb b/spec/has_aux_table_spec.rb
index c06def9..884603e 100644
--- a/spec/has_aux_table_spec.rb
+++ b/spec/has_aux_table_spec.rb
@@ -784,6 +784,117 @@ RSpec.describe HasAuxTable do
         expect(Driver.where(car: nil)).to eq([driver2])
         expect(Driver.where(car: nodriver_car)).to eq([])
       end
+
+      describe "custom foreign and primary keys" do
+        before(:each) do
+          @e1 =
+            ModelECustom.create!(
+              pk_base_id: 1,
+              fk_base_id: 2,
+              pk_aux_id: 3,
+              fk_aux_id: 4
+            )
+          @e2 =
+            ModelECustom.create!(
+              pk_base_id: 5,
+              fk_base_id: 6,
+              pk_aux_id: 7,
+              fk_aux_id: 8
+            )
+        end
+
+        describe "belongs_to association" do
+          it "works between base and base" do
+            # e1.fk_base_id <- e2.pk_base_id
+            @e1.base_to_base = @e2
+            @e1.save!
+            @e2.save!
+            expect(@e1.fk_base_id).to eq(5)
+            expect(@e2.pk_base_id).to eq(5)
+            expect(ModelECustom.where(base_to_base: @e2)).to eq([@e1])
+            expect(ModelECustom.where(base_to_base: @e1)).to eq([])
+          end
+
+          it "works between base and aux" do
+            # e1.fk_base_id <- e2.pk_aux_id
+            @e1.base_to_aux = @e2
+            @e1.save!
+            @e2.save!
+            expect(@e1.fk_base_id).to eq(7)
+            expect(@e2.pk_aux_id).to eq(7)
+            expect(ModelECustom.where(base_to_aux: @e2)).to eq([@e1])
+            expect(ModelECustom.where(base_to_aux: @e1)).to eq([])
+          end
+
+          it "works between aux and base" do
+            # e1.fk_aux_id <- e2.pk_base_id
+            @e1.aux_to_base = @e2
+            @e1.save!
+            @e2.save!
+            expect(@e1.fk_aux_id).to eq(5)
+            expect(@e2.pk_base_id).to eq(5)
+            expect(ModelECustom.where(aux_to_base: @e2)).to eq([@e1])
+            expect(ModelECustom.where(aux_to_base: @e1)).to eq([])
+          end
+
+          it "works between aux and aux" do
+            # e1.fk_aux_id <- e2.pk_aux_id
+            @e1.aux_to_aux = @e2
+            @e1.save!
+            @e2.save!
+            expect(@e1.fk_aux_id).to eq(7)
+            expect(@e2.pk_aux_id).to eq(7)
+            expect(ModelECustom.where(aux_to_aux: @e2)).to eq([@e1])
+            expect(ModelECustom.where(aux_to_aux: @e1)).to eq([])
+          end
+        end
+
+        describe "has_one association" do
+          it "works between base and base" do
+            # e1.pk_base_id -> e2.fk_base_id
+            @e1.owned_base_to_base = @e2
+            @e1.save!
+            @e2.save!
+            expect(@e1.pk_base_id).to eq(1)
+            expect(@e2.fk_base_id).to eq(1)
+            expect(ModelECustom.where(owned_base_to_base: @e2)).to eq([@e1])
+            expect(ModelECustom.where(owned_base_to_base: @e1)).to eq([])
+          end
+
+          it "works between base and aux" do
+            # e1.pk_base_id -> e2.fk_aux_id
+            @e1.owned_base_to_aux = @e2
+            @e1.save!
+            @e2.save!
+            expect(@e1.pk_base_id).to eq(1)
+            expect(@e2.fk_aux_id).to eq(1)
+            expect(ModelECustom.where(owned_base_to_aux: @e2)).to eq([@e1])
+            expect(ModelECustom.where(owned_base_to_aux: @e1)).to eq([])
+          end
+
+          it "works between aux and base" do
+            # e1.pk_aux_id -> e2.fk_base_id
+            @e1.owned_aux_to_base = @e2
+            @e1.save!
+            @e2.save!
+            expect(@e1.pk_aux_id).to eq(3)
+            expect(@e2.fk_base_id).to eq(3)
+            expect(ModelECustom.where(owned_aux_to_base: @e2)).to eq([@e1])
+            expect(ModelECustom.where(owned_aux_to_base: @e1)).to eq([])
+          end
+
+          it "works between aux and aux" do
+            # e1.pk_aux_id -> e2.fk_aux_id
+            @e1.owned_aux_to_aux = @e2
+            @e1.save!
+            @e2.save!
+            expect(@e1.pk_aux_id).to eq(3)
+            expect(@e2.fk_aux_id).to eq(3)
+            expect(ModelECustom.where(owned_aux_to_aux: @e2)).to eq([@e1])
+            expect(ModelECustom.where(owned_aux_to_aux: @e1)).to eq([])
+          end
+        end
+      end
     end
 
     describe "#reload" do
diff --git a/spec/spec_models.rb b/spec/spec_models.rb
index fb610a6..b0fbf9c 100644
--- a/spec/spec_models.rb
+++ b/spec/spec_models.rb
@@ -405,3 +405,64 @@ class AdJoin < ActiveRecord::Base
   belongs_to :model_a, inverse_of: :ad_joins, counter_cache: true
   belongs_to :model_d, inverse_of: :ad_joins, counter_cache: true
 end
+
+ActiveRecord::Schema.define do
+  create_base_table :model_es do |t|
+    t.integer :pk_base_id, index: true
+    t.integer :fk_base_id, index: true
+    t.create_aux :e do |t|
+      t.integer :pk_aux_id, index: true
+      t.integer :fk_aux_id, index: true
+    end
+  end
+end
+
+class ModelE < ActiveRecord::Base
+  include HasAuxTable
+end
+class ModelECustom < ModelE
+  aux_table :e
+  belongs_to :base_to_base,
+             class_name: "ModelECustom",
+             foreign_key: :fk_base_id,
+             primary_key: :pk_base_id,
+             optional: true
+
+  belongs_to :base_to_aux,
+             class_name: "ModelECustom",
+             foreign_key: :fk_base_id,
+             primary_key: :pk_aux_id,
+             optional: true
+
+  belongs_to :aux_to_base,
+             class_name: "ModelECustom",
+             foreign_key: :fk_aux_id,
+             primary_key: :pk_base_id,
+             optional: true
+
+  belongs_to :aux_to_aux,
+             class_name: "ModelECustom",
+             foreign_key: :fk_aux_id,
+             primary_key: :pk_aux_id,
+             optional: true
+
+  has_one :owned_base_to_base,
+          class_name: "ModelECustom",
+          primary_key: :pk_base_id,
+          foreign_key: :fk_base_id
+
+  has_one :owned_base_to_aux,
+          class_name: "ModelECustom",
+          primary_key: :pk_base_id,
+          foreign_key: :fk_aux_id
+
+  has_one :owned_aux_to_base,
+          class_name: "ModelECustom",
+          primary_key: :pk_aux_id,
+          foreign_key: :fk_base_id
+
+  has_one :owned_aux_to_aux,
+          class_name: "ModelECustom",
+          primary_key: :pk_aux_id,
+          foreign_key: :fk_aux_id
+end