1 files changed, 29 insertions, 7 deletions
diff --git a/spec/workers/stuck_ci_jobs_worker_spec.rb b/spec/workers/stuck_ci_jobs_worker_spec.rb
index bdc64c6785b..856886e3df5 100644
--- a/spec/workers/stuck_ci_jobs_worker_spec.rb
+++ b/spec/workers/stuck_ci_jobs_worker_spec.rb
@@ -1,14 +1,21 @@
 require 'spec_helper'
 
 describe StuckCiJobsWorker do
+  include ExclusiveLeaseHelpers
+
   let!(:runner) { create :ci_runner }
   let!(:job) { create :ci_build, runner: runner }
-  let(:worker) { described_class.new }
-  let(:exclusive_lease_uuid) { SecureRandom.uuid }
+  let(:trace_lease_key) { "trace:archive:#{job.id}" }
+  let(:trace_lease_uuid) { SecureRandom.uuid }
+  let(:worker_lease_key) { StuckCiJobsWorker::EXCLUSIVE_LEASE_KEY }
+  let(:worker_lease_uuid) { SecureRandom.uuid }
+
+  subject(:worker) { described_class.new }
 
   before do
+    stub_exclusive_lease(worker_lease_key, worker_lease_uuid)
+    stub_exclusive_lease(trace_lease_key, trace_lease_uuid)
     job.update!(status: status, updated_at: updated_at)
-    allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain).and_return(exclusive_lease_uuid)
   end
 
   shared_examples 'job is dropped' do
@@ -44,16 +51,19 @@ describe StuckCiJobsWorker do
 
       context 'when job was not updated for more than 1 day ago' do
         let(:updated_at) { 2.days.ago }
+
         it_behaves_like 'job is dropped'
       end
 
       context 'when job was updated in less than 1 day ago' do
         let(:updated_at) { 6.hours.ago }
+
         it_behaves_like 'job is unchanged'
       end
 
       context 'when job was not updated for more than 1 hour ago' do
         let(:updated_at) { 2.hours.ago }
+
         it_behaves_like 'job is unchanged'
       end
     end
@@ -65,11 +75,14 @@ describe StuckCiJobsWorker do
 
       context 'when job was not updated for more than 1 hour ago' do
         let(:updated_at) { 2.hours.ago }
+
         it_behaves_like 'job is dropped'
       end
 
-      context 'when job was updated in less than 1 hour ago' do
+      context 'when job was updated in less than 1
+       hour ago' do
         let(:updated_at) { 30.minutes.ago }
+
         it_behaves_like 'job is unchanged'
       end
     end
@@ -80,11 +93,13 @@ describe StuckCiJobsWorker do
 
     context 'when job was not updated for more than 1 hour ago' do
       let(:updated_at) { 2.hours.ago }
+
       it_behaves_like 'job is dropped'
     end
 
     context 'when job was updated in less than 1 hour ago' do
       let(:updated_at) { 30.minutes.ago }
+
       it_behaves_like 'job is unchanged'
     end
   end
@@ -93,6 +108,7 @@ describe StuckCiJobsWorker do
     context "when job is #{status}" do
       let(:status) { status }
       let(:updated_at) { 2.days.ago }
+
       it_behaves_like 'job is unchanged'
     end
   end
@@ -119,20 +135,26 @@ describe StuckCiJobsWorker do
     it 'is guard by exclusive lease when executed concurrently' do
       expect(worker).to receive(:drop).at_least(:once).and_call_original
       expect(worker2).not_to receive(:drop)
+
       worker.perform
-      allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain).and_return(false)
+
+      stub_exclusive_lease_taken(worker_lease_key)
+
       worker2.perform
     end
 
     it 'can be executed in sequence' do
       expect(worker).to receive(:drop).at_least(:once).and_call_original
       expect(worker2).to receive(:drop).at_least(:once).and_call_original
+
       worker.perform
       worker2.perform
     end
 
-    it 'cancels exclusive lease after worker perform' do
-      expect(Gitlab::ExclusiveLease).to receive(:cancel).with(described_class::EXCLUSIVE_LEASE_KEY, exclusive_lease_uuid)
+    it 'cancels exclusive leases after worker perform' do
+      expect_to_cancel_exclusive_lease(trace_lease_key, trace_lease_uuid)
+      expect_to_cancel_exclusive_lease(worker_lease_key, worker_lease_uuid)
+
       worker.perform
     end
   end