tests for has_one & belongs_to custom foreign keys

lib/has_aux_table/aux_table_config.rb

@@ -141,8 +141,9 @@ module HasAuxTable
           # attribute is a column on the aux table
           aux[k] = v
         elsif assoc = self.main.klass.reflect_on_association(k.to_s)
+          if assoc.is_a?(ActiveRecord::Reflection::BelongsToReflection)
             # attribute is an association on the main class
-          fk = assoc.association_foreign_key
+            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
@@ -151,6 +152,16 @@ module HasAuxTable
               # 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
+            main[k] = v
+          end
         else
           # attribute is not a column on the aux table or an association,
           # assume it's a column on the main table

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

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