Allow custom stamper columns to be respected

lowjoel · Feb 9, 2017 · 85faa46 · 85faa46
1 parent efc2b36
commit 85faa46
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 11 deletions.
diff --git a/lib/active_record/userstamp/stampable.rb b/lib/active_record/userstamp/stampable.rb
@@ -102,7 +102,7 @@ def set_creator_attribute
     current_attribute_value = send(attribute)
     return if current_attribute_value.present?
 
-    ActiveRecord::Userstamp::Utilities.assign_attribute(self, attribute)
+    ActiveRecord::Userstamp::Utilities.assign_stamper(self, :creator, attribute)
   end
 
   def set_updater_attribute
@@ -111,14 +111,14 @@ def set_updater_attribute
 
     return if !self.new_record? && !self.changed?
 
-    ActiveRecord::Userstamp::Utilities.assign_attribute(self, attribute)
+    ActiveRecord::Userstamp::Utilities.assign_stamper(self, :updater, attribute)
   end
 
   def set_deleter_attribute
     attribute = ActiveRecord::Userstamp.config.deleter_attribute
     return if !has_stamper? || attribute.nil? || !has_attribute?(attribute)
 
-    ActiveRecord::Userstamp::Utilities.assign_attribute(self, attribute)
+    ActiveRecord::Userstamp::Utilities.assign_stamper(self, :deleter, attribute)
     save
   end
 end
diff --git a/lib/active_record/userstamp/utilities.rb b/lib/active_record/userstamp/utilities.rb
@@ -38,18 +38,17 @@ def self.available_association_columns(model)
     nil
   end
 
-  # Assigns the given attribute to the record, based on the model's stamper.
+  # Assigns the stamper to the given association/attribute in the record.
   #
-  # If the stamper is a record, then it is assigned to the attribute; if it is a number, then it
-  # is assigned to the +_id+ attribute
+  # If the stamper is a record, then it is assigned to the association; if it is a number, then it
+  # is assigned to the configured attribute.
   #
   # @param [ActiveRecord::Base] record The record to assign.
-  # @param [Symbol] attribute The attribute to assign.
-  def self.assign_attribute(record, attribute)
-    attribute = attribute.to_s
+  # @param [Symbol] association The association to assign, if the stamper is a record.
+  # @param [Symbol] attribute The attribute to assign, if the stamper is an ID.
+  def self.assign_stamper(record, association, attribute)
     stamp_value = record.class.stamper_class.stamper
-
-    attribute = attribute[0..-4] if !stamp_value.is_a?(Fixnum) && attribute.end_with?('_id')
+    attribute = association if !stamp_value.is_a?(Fixnum) && !stamp_value.is_a?(Integer)
     record.send("#{attribute}=", stamp_value)
   end
 end