Add helper methods to stub Gitlab::ExclusiveLease

3 files changed, 42 insertions, 6 deletions

diff --git a/spec/support/helpers/exclusive_lease_helpers.rb b/spec/support/helpers/exclusive_lease_helpers.rb
new file mode 100644
index 00000000000..383cc7dee81
--- /dev/null
+++ b/spec/support/helpers/exclusive_lease_helpers.rb
@@ -0,0 +1,36 @@
+module ExclusiveLeaseHelpers
+  def stub_exclusive_lease(key = nil, uuid = 'uuid', renew: false, timeout: nil)
+    key     ||= instance_of(String)
+    timeout ||= instance_of(Integer)
+
+    lease = instance_double(
+      Gitlab::ExclusiveLease,
+      try_obtain: uuid,
+      exists?: true,
+      renew: renew
+    )
+
+    allow(Gitlab::ExclusiveLease)
+      .to receive(:new)
+      .with(key, timeout: timeout)
+      .and_return(lease)
+
+    lease
+  end
+
+  def stub_exclusive_lease_taken(key = nil, timeout: nil)
+    stub_exclusive_lease(key, nil, timeout: timeout)
+  end
+
+  def expect_to_obtain_exclusive_lease(key, uuid = 'uuid', timeout: nil)
+    lease = stub_exclusive_lease(key, uuid, timeout: timeout)
+
+    expect(lease).to receive(:try_obtain)
+  end
+
+  def expect_to_cancel_exclusive_lease(key, uuid)
+    expect(Gitlab::ExclusiveLease)
+      .to receive(:cancel)
+      .with(key, uuid)
+  end
+end
diff --git a/spec/support/shared_examples/ci_trace_shared_examples.rb b/spec/support/shared_examples/ci_trace_shared_examples.rb
index 6dbe0f6f980..db723a323f8 100644
--- a/spec/support/shared_examples/ci_trace_shared_examples.rb
+++ b/spec/support/shared_examples/ci_trace_shared_examples.rb
@@ -247,8 +247,10 @@ shared_examples_for 'common trace features' do
         end
 
         context 'when another process has already been archiving', :clean_gitlab_redis_shared_state do
+          include ExclusiveLeaseHelpers
+
           before do
-            Gitlab::ExclusiveLease.new("trace:archive:#{trace.job.id}", timeout: 1.hour).try_obtain
+            stub_exclusive_lease_taken("trace:archive:#{trace.job.id}", timeout: 1.hour)
           end
 
           it 'blocks concurrent archiving' do
diff --git a/spec/support/shared_examples/uploaders/object_storage_shared_examples.rb b/spec/support/shared_examples/uploaders/object_storage_shared_examples.rb
index 19800c6638f..1bd176280c5 100644
--- a/spec/support/shared_examples/uploaders/object_storage_shared_examples.rb
+++ b/spec/support/shared_examples/uploaders/object_storage_shared_examples.rb
@@ -76,8 +76,10 @@ shared_examples "migrates" do |to_store:, from_store: nil|
   end
 
   context 'when migrate! is occupied by another process' do
+    include ExclusiveLeaseHelpers
+
     before do
-      @uuid = Gitlab::ExclusiveLease.new(subject.exclusive_lease_key, timeout: 1.hour.to_i).try_obtain
+      stub_exclusive_lease_taken(subject.exclusive_lease_key, timeout: 1.hour.to_i)
     end
 
     it 'does not execute migrate!' do
@@ -91,10 +93,6 @@ shared_examples "migrates" do |to_store:, from_store: nil|
 
       expect { subject.use_file }.to raise_error(ObjectStorage::ExclusiveLeaseTaken)
     end
-
-    after do
-      Gitlab::ExclusiveLease.cancel(subject.exclusive_lease_key, @uuid)
-    end
   end
 
   context 'migration is unsuccessful' do