diff --git a/lib/mongo_mapper/plugins/userstamps.rb b/lib/mongo_mapper/plugins/userstamps.rb index 291620b34..fc9c5ebf0 100644 --- a/lib/mongo_mapper/plugins/userstamps.rb +++ b/lib/mongo_mapper/plugins/userstamps.rb @@ -5,11 +5,14 @@ module Userstamps extend ActiveSupport::Concern module ClassMethods - def userstamps! + def userstamps!(class_name = 'User') key :creator_id, ObjectId key :updater_id, ObjectId - belongs_to :creator, :class_name => 'User' - belongs_to :updater, :class_name => 'User' + belongs_to :creator, :class_name => class_name + belongs_to :updater, :class_name => class_name + end + def userstamps_for!(class_name = 'User') + userstamps!(class_name) end end end diff --git a/test/functional/test_userstamps.rb b/test/functional/test_userstamps.rb index 94bcd40af..c00f58e06 100644 --- a/test/functional/test_userstamps.rb +++ b/test/functional/test_userstamps.rb @@ -6,22 +6,28 @@ class UserstampsTest < Test::Unit::TestCase @document = Doc do userstamps! end + @document_alt_user = Doc do + userstamps! 'AltUser' + end + @document_for_alt_user = Doc do + userstamps_for! 'AltUser' + end end should "add creator_id key" do - @document.keys.keys.should include('creator_id') + [@document, @document_alt_user, @document_for_alt_user].each{ |d| d.keys.keys.should include('creator_id') } end should "add updater_id key" do - @document.keys.keys.should include('updater_id') + [@document, @document_alt_user, @document_for_alt_user].each{ |d| d.keys.keys.should include('updater_id') } end should "add belongs_to creator" do - @document.associations.keys.should include(:creator) + [@document, @document_alt_user, @document_for_alt_user].each{ |d| d.associations.keys.should include(:creator) } end should "add belongs_to updater" do - @document.associations.keys.should include(:updater) + [@document, @document_alt_user, @document_for_alt_user].each{ |d| d.associations.keys.should include(:updater) } end end end \ No newline at end of file