lib/ensurance.rb in ensurance-0.1.2 vs lib/ensurance.rb in ensurance-0.1.3

@@ -8,12 +8,12 @@
 module Ensurance
   extend ActiveSupport::Concern
 
   class_methods do
     def ensure_by(*args)
-      @_ensure_by ||= [:id]
-      @_ensure_by += [args].flatten
+      @_additional_ensure_by = args
+      @_ensure_by = nil
     end
 
     def ensure(thing = nil)
       return nil unless thing.present?
 
@@ -26,10 +26,10 @@
       elsif thing.is_a?(String) && found = GlobalID::Locator.locate(thing)
         return found
       end
       found = nil
 
-      @_ensure_by ||= [:id]
+      @_ensure_by ||= [self.primary_key, @_additional_ensure_by].flatten.compact.uniq
       @_ensure_by.each do |ensure_field|
         value = thing.try(:fetch, ensure_field.to_sym, nil) || thing.try(:fetch, ensure_field.to_s, nil) || thing
         found = self.find_by(ensure_field => value)
         break if found.is_a?(self)
       end