Add with_attributes to override both transaction and organization ids

lib/audited/audit.rb

@@ -13,6 +13,7 @@ def setup_audit
 
         before_create :set_audit_user
         before_create :set_transaction_id
+        before_create :set_attributes
 
         cattr_accessor :audited_class_names
         self.audited_class_names = Set.new
@@ -45,6 +46,18 @@ def with_transaction_id(transaction_id, &block)
         yieldval
       end
 
+      # Override a bunch of values
+      def with_attributes(hash, &block)
+        hash.each do |key, value|
+          Thread.current[key] = value
+        end
+        yieldval = yield
+        hash.each do |key, value|
+          Thread.current[key] = nil
+        end
+        yieldval
+      end
+
     end
 
     # Returns a hash of the changed attributes with the new values
@@ -70,6 +83,12 @@ def set_audit_user
       nil # prevent stopping callback chains
     end
 
+    def set_attributes
+      self.transaction_id = Thread.current[:audited_transaction_id] if Thread.current[:audited_transaction_id]
+      self.organization_id = Thread.current[:audited_organization_id] if Thread.current[:audited_organization_id]
+      nil
+    end
+
     def set_transaction_id
       self.transaction_id = Thread.current[:audited_transaction_id] if Thread.current[:audited_transaction_id]
       nil

spec/audited/adapters/active_record/audit_spec.rb

@@ -106,7 +106,34 @@ class Models::ActiveRecord::CustomUserSubclass < Models::ActiveRecord::CustomUse
         tr
       end.should == tr
     end
+  end
+
+  describe "with_attributes" do
+    let(:tr) { SecureRandom.hex(32) }
+    let(:org) { 123 }
+
+    it "should record transaction ids" do
+      attributes = {
+        :audited_transaction_id => tr,
+        :audited_organization_id => org
+      }
+      Audited.audit_class.with_attributes(attributes) do
+        company = Models::ActiveRecord::Company.create :name => 'The auditors'
+        company.name = 'The Auditors, Inc'
+        company.save
 
+        company.audits.each do |audit|
+          audit.transaction_id.should == tr
+          audit.organization_id.should == org
+        end
+      end
+    end
+
+    it "should return the value from the yield block" do
+      Audited.audit_class.with_transaction_id(tr) do
+        tr
+      end.should == tr
+    end
   end
 
   describe "as_user" do