From a9f4c24b4ecfed4deb833ae9301c49befd8a14c1 Mon Sep 17 00:00:00 2001
From: Overload119 <budgeneration@gmail.com>
Date: Fri, 26 Apr 2013 11:16:18 -0400
Subject: [PATCH] Add with_attributes to override both transaction and
 organization ids

---
 lib/audited/audit.rb                          | 19 +++++++++++++
 .../adapters/active_record/audit_spec.rb      | 27 +++++++++++++++++++
 2 files changed, 46 insertions(+)

diff --git a/lib/audited/audit.rb b/lib/audited/audit.rb
index 8bc614ab1..3e9b9ddb6 100644
--- a/lib/audited/audit.rb
+++ b/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
diff --git a/spec/audited/adapters/active_record/audit_spec.rb b/spec/audited/adapters/active_record/audit_spec.rb
index 89ce81fdd..402d5ad40 100644
--- a/spec/audited/adapters/active_record/audit_spec.rb
+++ b/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