write attribute with polymorphic integer type

Author: wendy-clio

lib/polymorphic_integer_type.rb

@@ -6,6 +6,7 @@
 require "polymorphic_integer_type/module_generator"
 require "polymorphic_integer_type/belongs_to_polymorphic_association_extension"
 require "polymorphic_integer_type/activerecord_5_0_0/polymorphic_array_value_extension"
+require "polymorphic_integer_type/polymorphic_foreign_association_extension"
 
 if ACTIVE_RECORD_VERSION < Gem::Version.new("5.2.0")
   require "polymorphic_integer_type/activerecord_5_0_0/association_query_handler_extension"

lib/polymorphic_integer_type/extensions.rb

@@ -58,6 +58,8 @@ def remove_type_and_establish_mapping(name, options, scope)
 
           options[:foreign_key] ||= "#{poly_type}_id"
           foreign_type = options.delete(:foreign_type) || "#{poly_type}_type"
+          options[:foreign_integer_type] = foreign_type
+          options[:integer_type] = klass_mapping.to_i
 
           options[:scope] ||= -> {
             condition = where(foreign_type => klass_mapping.to_i)
@@ -87,7 +89,9 @@ def has_many(name, scope = nil, **options, &extension)
         end
 
         remove_type_and_establish_mapping(name, options, scope)
-        super(name, options.delete(:scope), **options, &extension)
+        super(name, options.delete(:scope), **options.except(:foreign_integer_type, :integer_type), &extension).tap do |_|
+          remove_integer_type_and_set_attributes_and_extension(options, ActiveRecord::Reflection::HasManyReflection, reflections[name.to_s])
+        end
       end
 
       def has_one(name, scope = nil, **options)
@@ -97,7 +101,28 @@ def has_one(name, scope = nil, **options)
         end
 
         remove_type_and_establish_mapping(name, options, scope)
-        super(name, options.delete(:scope), **options)
+        super(name, options.delete(:scope), **options.except(:foreign_integer_type, :integer_type)).tap do |_|
+          remove_integer_type_and_set_attributes_and_extension(options, ActiveRecord::Reflection::HasOneReflection, reflections[name.to_s])
+        end
+      end
+
+      def remove_integer_type_and_set_attributes_and_extension(options, klass, reflection)
+        foreign_integer_type = options.delete :foreign_integer_type
+        integer_type = options.delete :integer_type
+        is_polymorphic_integer = foreign_integer_type && integer_type
+
+        if is_polymorphic_integer
+          klass.attr_accessor(:foreign_integer_type)
+          klass.attr_accessor(:integer_type)
+          reflection.foreign_integer_type = foreign_integer_type
+          reflection.integer_type = integer_type
+
+          if Gem::Version.new(ActiveRecord::VERSION::STRING) < Gem::Version.new("6.1")
+            ActiveRecord::Associations::Association.prepend(ActiveRecord::Associations::PolymorphicForeignAssociationExtension)
+          else
+            ActiveRecord::Associations::ForeignAssociation.prepend(ActiveRecord::Associations::PolymorphicForeignAssociationExtension)
+          end
+        end
       end
 
 

lib/polymorphic_integer_type/polymorphic_foreign_association_extension.rb

@@ -0,0 +1,13 @@
+module ActiveRecord
+  module Associations
+    module PolymorphicForeignAssociationExtension
+
+      def set_owner_attributes(record)
+        super
+        if reflection.foreign_integer_type && reflection.integer_type
+          record._write_attribute(reflection.foreign_integer_type, reflection.integer_type)
+        end
+      end
+    end
+  end
+end

spec/polymorphic_integer_type_spec.rb

@@ -26,6 +26,28 @@
       expect(link.target_type).to eq("Food")
     end
 
+    context "from HasManyReflection" do
+      it "sets the source properly from HasManyReflection" do
+        link_1 = Link.create()
+        link_2 = Link.create()
+        dog.source_links = [link_1, link_2]
+        expect(link_1.source_type).to eq("Animal")
+        expect(link_1.source_id).to eq(dog.id)
+        expect(link_2.source_type).to eq("Animal")
+        expect(link_1.source_id).to eq(dog.id)
+      end
+    end
+
+    context "from HasOneReflection" do
+      it "sets the source properly from HasManyReflection" do
+        link = Link.create()
+        dog.source_link = link
+
+        expect(link.source_type).to eq("Animal")
+        expect(link.source_id).to eq(dog.id)
+      end
+    end
+
     context "when models are namespaced" do
       context "and mappings include namespaces" do
         it "sets the source_type" do
@@ -340,6 +362,10 @@ class InlineDrink2 < ActiveRecord::Base
       expect(link[:target_type]).to eq(13)
     end
 
+    it "pulls mapping from given hash" do
+      animal.source_links.new
+    end
+
     it "doesn't break string type polymorphic associations" do
       expect(link.normal_target).to eq(drink)
       expect(link.normal_target_type).to eq("InlineDrink2")

spec/support/animal.rb

@@ -3,5 +3,5 @@ class Animal < ActiveRecord::Base
 
   belongs_to :owner, class_name: "Person"
   has_many :source_links, as: :source, integer_type: true, class_name: "Link"
-
+  has_one :source_link, as: :source, integer_type: true, class_name: "Link"
 end