From e69a2606cd79f9bd015a5dbd0d3d1481ac52c46c Mon Sep 17 00:00:00 2001
From: Joel Low <joel@joelsplace.sg>
Date: Thu, 9 Feb 2017 10:21:48 +0800
Subject: [PATCH] Simplify the attribute assignment logic

This also fixes model-only configuration for the stamp columns
---
 lib/active_record/userstamp/stampable.rb | 27 +++++++++++++-----------
 lib/active_record/userstamp/utilities.rb | 17 +++++++++------
 2 files changed, 26 insertions(+), 18 deletions(-)

diff --git a/lib/active_record/userstamp/stampable.rb b/lib/active_record/userstamp/stampable.rb
index 83fe49e..2fcb290 100644
--- a/lib/active_record/userstamp/stampable.rb
+++ b/lib/active_record/userstamp/stampable.rb
@@ -96,29 +96,32 @@ def has_stamper?
   end
 
   def set_creator_attribute
-    attribute = ActiveRecord::Userstamp.config.creator_attribute
-    return if !has_stamper? || attribute.nil? || !has_attribute?(attribute)
+    return unless has_stamper?
 
-    current_attribute_value = send(attribute)
-    return if current_attribute_value.present?
+    creator_association = self.class.reflect_on_association(:creator)
+    return unless creator_association
+    return if creator.present?
 
-    ActiveRecord::Userstamp::Utilities.assign_stamper(self, :creator, attribute)
+    ActiveRecord::Userstamp::Utilities.assign_stamper(self, creator_association)
   end
 
   def set_updater_attribute
-    attribute = ActiveRecord::Userstamp.config.updater_attribute
-    return if !has_stamper? || attribute.nil? || !has_attribute?(attribute)
+    return unless has_stamper?
 
-    return if !self.new_record? && !self.changed?
+    updater_association = self.class.reflect_on_association(:updater)
+    return unless updater_association
+    return if !new_record? && !changed?
 
-    ActiveRecord::Userstamp::Utilities.assign_stamper(self, :updater, attribute)
+    ActiveRecord::Userstamp::Utilities.assign_stamper(self, updater_association)
   end
 
   def set_deleter_attribute
-    attribute = ActiveRecord::Userstamp.config.deleter_attribute
-    return if !has_stamper? || attribute.nil? || !has_attribute?(attribute)
+    return unless has_stamper?
 
-    ActiveRecord::Userstamp::Utilities.assign_stamper(self, :deleter, attribute)
+    deleter_association = self.class.reflect_on_association(:deleter)
+    return unless deleter_association
+
+    ActiveRecord::Userstamp::Utilities.assign_stamper(self, deleter_association)
     save
   end
 end
diff --git a/lib/active_record/userstamp/utilities.rb b/lib/active_record/userstamp/utilities.rb
index 7223250..4586b9e 100644
--- a/lib/active_record/userstamp/utilities.rb
+++ b/lib/active_record/userstamp/utilities.rb
@@ -38,17 +38,22 @@ def self.available_association_columns(model)
     nil
   end
 
-  # Assigns the stamper to the given association/attribute in the record.
+  # Assigns the stamper to the given association reflection in the record.
   #
   # 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.
+  # is assigned to the foreign key.
   #
   # @param [ActiveRecord::Base] record The record to assign.
-  # @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)
+  # @param [ActiveRecord::Reflection] association The association to assign
+  def self.assign_stamper(record, association)
     stamp_value = record.class.stamper_class.stamper
-    attribute = association if !stamp_value.is_a?(Fixnum) && !stamp_value.is_a?(Integer)
+    attribute =
+      if stamp_value.is_a?(ActiveRecord::Base)
+        association.name
+      else
+        association.foreign_key
+      end
+
     record.send("#{attribute}=", stamp_value)
   end
 end