From ec8d88dda3bf2ccec12869dc3173b054d5296655 Mon Sep 17 00:00:00 2001
From: tvongaza <tys.von.gaza@aeon.health>
Date: Fri, 22 Sep 2023 13:10:40 -0600
Subject: [PATCH] Add a new ActiveJobExtensions which hooks into the
 serialization and deserialization of an ActiveJob to preserve the current
 tenant.

---
 lib/acts_as_tenant.rb                       |  5 +++
 lib/acts_as_tenant/active_job_extensions.rb | 13 ++++++
 spec/jobs/active_job_extensions_spec.rb     | 49 +++++++++++++++++++++
 3 files changed, 67 insertions(+)
 create mode 100644 lib/acts_as_tenant/active_job_extensions.rb
 create mode 100644 spec/jobs/active_job_extensions_spec.rb

diff --git a/lib/acts_as_tenant.rb b/lib/acts_as_tenant.rb
index ca4160f..6d31617 100644
--- a/lib/acts_as_tenant.rb
+++ b/lib/acts_as_tenant.rb
@@ -8,6 +8,7 @@ module ActsAsTenant
   autoload :ControllerExtensions, "acts_as_tenant/controller_extensions"
   autoload :ModelExtensions, "acts_as_tenant/model_extensions"
   autoload :TenantHelper, "acts_as_tenant/tenant_helper"
+  autoload :ActiveJobExtensions, "acts_as_tenant/active_job_extensions"
 
   @@configuration = nil
   @@tenant_klass = nil
@@ -158,3 +159,7 @@ def self.should_require_tenant?
 ActiveSupport.on_load(:action_view) do |base|
   base.include ActsAsTenant::TenantHelper
 end
+
+ActiveSupport.on_load(:active_job) do |base|
+  base.prepend ActsAsTenant::ActiveJobExtensions
+end
diff --git a/lib/acts_as_tenant/active_job_extensions.rb b/lib/acts_as_tenant/active_job_extensions.rb
new file mode 100644
index 0000000..b9b5d96
--- /dev/null
+++ b/lib/acts_as_tenant/active_job_extensions.rb
@@ -0,0 +1,13 @@
+module ActsAsTenant
+  module ActiveJobExtensions
+    def serialize
+      super.merge("current_tenant" => ActsAsTenant.current_tenant&.to_global_id)
+    end
+
+    def deserialize(job_data)
+      tenant_global_id = job_data.delete("current_tenant")
+      ActsAsTenant.current_tenant = tenant_global_id ? GlobalID::Locator.locate(tenant_global_id) : nil
+      super
+    end
+  end
+end
diff --git a/spec/jobs/active_job_extensions_spec.rb b/spec/jobs/active_job_extensions_spec.rb
new file mode 100644
index 0000000..548f205
--- /dev/null
+++ b/spec/jobs/active_job_extensions_spec.rb
@@ -0,0 +1,49 @@
+require "spec_helper"
+
+class ApplicationTestJob < ApplicationJob
+  def perform(expected_tenant:)
+    raise ApplicationTestJobTenantError unless ActsAsTenant.current_tenant == expected_tenant
+    Project.all
+  end
+end
+
+class ApplicationTestJobTenantError < StandardError; end
+
+RSpec.describe ApplicationTestJob, type: :job do
+  include ActiveJob::TestHelper
+
+  let(:account) { accounts(:foo) }
+
+  describe "#perform_later" do
+    context "when tenant is required" do
+      before { allow(ActsAsTenant.configuration).to receive_messages(require_tenant: true) }
+
+      it "raises ApplicationTestJobTenantError when expected_tenant does not match current_tenant" do
+        ActsAsTenant.current_tenant = account
+        expect { described_class.perform_later(expected_tenant: nil) }.to have_enqueued_job.on_queue("default")
+        expect { perform_enqueued_jobs }.to raise_error(ApplicationTestJobTenantError)
+      end
+
+      it "when tenant is set, successfully queues and performs job" do
+        ActsAsTenant.current_tenant = account
+        expect { described_class.perform_later(expected_tenant: account) }.to have_enqueued_job.on_queue("default")
+        expect { perform_enqueued_jobs }.not_to raise_error
+      end
+
+      it "when tenant is not set, successfully queues but fails to perform job" do
+        ActsAsTenant.current_tenant = nil
+        expect { described_class.perform_later(expected_tenant: nil) }.to have_enqueued_job.on_queue("default")
+        expect { perform_enqueued_jobs }.to raise_error(ActsAsTenant::Errors::NoTenantSet)
+      end
+    end
+
+    context "when tenant is not required" do
+      before { allow(ActsAsTenant.configuration).to receive_messages(require_tenant: false) }
+      it "when tenant is not set, queues and performs job" do
+        ActsAsTenant.current_tenant = nil
+        expect { described_class.perform_later(expected_tenant: nil) }.to have_enqueued_job.on_queue("default")
+        expect { perform_enqueued_jobs }.not_to raise_error
+      end
+    end
+  end
+end